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

我想在 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 

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

到目前为止，我在 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 的功能先进，但这在某种程度上仍然是可能的。

请您参考如下方法:

您尝试做的事情经常用 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]]