TensorFlow 使用数组索引将 Tensor 分配给 Tensor

Question

我想在 TensorFlow 中做类似这段 Numpy 代码的事情：

a = np.zeros([5, 2])
idx = np.random.randint(0, 2, (5,))
row_idx = np.arange(5)
a[row_idx, idx] = row_idx

意思是用另一个张量索引一个 2D 张量的所有行，然后为其分配一个张量。我对如何实现这一目标完全一无所知。

到目前为止，我在 Tensorflow 中可以做的如下

a = tf.Variable(tf.zeros((5, 2)))
idx = tf.constant([0, 1, 1, 0, 1])
row_idx = tf.range(5)
indices = tf.transpose([row_idx, idx])
r = tf.gather_nd(a, indices)
tf.assign(r, row_idx) # This line does not work

当我尝试执行此操作时，在最后一行中出现以下错误：

AttributeError: 'Tensor' object has no attribute 'assign'

有没有解决的办法？必须有一些很好的方法来做到这一点，我不想用 for 循环遍历数据并在每个元素的基础上手动分配它。我知道现在数组索引不像 Numpy 的功能那么先进，但这应该仍然是可能的。

Answer 1

你想要做的事情经常是用 完成的tf.scatter_nd_update。然而，大多数时候这不是正确的方法，您不需要变量，只需要从原始张量生成的另一个张量和一些替换值。不幸的是，一般来说没有直接的方法可以做到这一点。如果您的原始张量确实全为零，那么您可以简单地使用tf.scatter_nd：

import tensorflow as tf

idx = tf.constant([0, 1, 1, 0, 1])
row_idx = tf.range(5)
indices = tf.stack([row_idx, idx], axis=1)
a = tf.scatter_nd(indices, row_idx, (5, 2))
with tf.Session() as sess:
    print(sess.run(a))
# [[0 0]
#  [0 1]
#  [0 2]
#  [3 0]
#  [0 4]]

然而，如果初始张量不全为零，则情况会更复杂。一种方法是执行与上面相同的操作，然后为更新的内容创建一个掩码，并根据掩码在原始和更新之间进行选择：

import tensorflow as tf

a = tf.ones((5, 2), dtype=tf.int32)
idx = tf.constant([0, 1, 1, 0, 1])
row_idx = tf.range(5)
indices = tf.stack([row_idx, idx], axis=1)
a_update = tf.scatter_nd(indices, row_idx, (5, 2))
update_mask = tf.scatter_nd(indices, tf.ones_like(row_idx, dtype=tf.bool), (5, 2))
a = tf.where(update_mask, a_update, a)
with tf.Session() as sess:
    print(sess.run(a))
# [[0 1]
#  [1 1]
#  [1 2]
#  [3 1]
#  [1 4]]