Tensorflow：如何用 "multiples" 张量平铺张量的元素？

Question

假设，我有一个输入张量如下：

a = tf.constant ([[0, 0], [1, 1], [2, 2]])

和一个"multiple"张量：

mul= tf.constant([1, 3, 2])

我想要这样的结果张量：

res =
 [[0, 0],
  [1, 1],
  [1, 1],
  [1, 1],
  [2, 2],
  [2, 2]]

张量a的行数(n)等于第张量 mul 中的元素。如果n是固定的，我可以使用下面的代码：

res = tf.tile ([a[0]], [mul[0], 1])

for i in range (1, 3):
    res = tf.concat ((res, tf.tile ([a[i]], [mul[i], 1])), 0)

但是如果我不知道 n（它会变化），我怎样才能得到结果？

如果您有任何想法，我将不胜感激！

Answer 1

你最简单的选择可能是使用 py_func 将 [1, 3, 2] 变成类似 [0, 1, 1, 1, 2, 2] 的东西，它可以用作索引tf.gather.

如果静态知道 a 张量中的最大行数，我想你可以使用 tf.dynamic_partition 来获取它，但它可能会导致一个大图。

或者，tf.while_loop 加上 tf.TensorArray 来存储输出也可以。

Tensorflow: How to tile elements of a Tensor with a "multiples" tensor?