python - TensorFlow 获取特定列的每一行的元素

Question

如果 A 是像这样的 TensorFlow 变量

A = tf.Variable([[1, 2], [3, 4]])

和index是另一个变量

index = tf.Variable([0, 1])

我想使用这个索引来选择每行中的列。在这种情况下，第一行的项目 0 和第二行的项目 1。

如果 A 是一个 Numpy 数组，那么要获取索引中提到的相应行的列，我们可以这样做

x = A[np.arange(A.shape[0]), index]

结果是

[1, 4]

TensorFlow 的等效操作是什么？我知道 TensorFlow 不支持很多索引操作。如果不能直接完成，有什么解决办法？

Answer 1

您可以使用行索引扩展列索引，然后使用 gather_nd:

import tensorflow as tf

A = tf.constant([[1, 2], [3, 4]])
indices = tf.constant([1, 0])

# prepare row indices
row_indices = tf.range(tf.shape(indices)[0])

# zip row indices with column indices
full_indices = tf.stack([row_indices, indices], axis=1)

# retrieve values by indices
S = tf.gather_nd(A, full_indices)

session = tf.InteractiveSession()
session.run(S)