tensorflow - 如何移动张量中的值

Question

我有带值的形状 [batch_size, A] 的张量 T 和带移位参数的形状 [batch_size] 的张量 S。

我想将 T[b] 中的值向右移动 S[b] 个位置，应删除 T[b] 的最后一个 S[b] 元素，并将新元素设置为 0。

所以基本上想做这样的事情:

for i in range(batch_size):
  T[i] = zeros[:S[i]] + T[i, :A-S[i]]

例子:

For:
T = [[1, 2, 3], [4, 5, 6]]
S = [1, 2]

Return:
T' = [[0, 1, 2], [0, 0, 4]]

有什么简单的方法吗？

Answer 1

您可以为此目的使用 tf.concat 和 tf.stack:

T_shift = tf.zeros((batch_size, A), tf.float32)
tmp = []

for i in xrange(batch_size):
    tmp.append(tf.concat([T_shift[i, :S[i, 0]],T[i, :17 - S[i,0]]], axis = 0))
T_shift = tf.stack(tmp)