Tensorflow：每行索引

Question

假设我有一个形状为 (100,20) 的张量。现在我还有一个形状指数的张量（100，）。现在如何获得形状为 (100,) 或 (100,1) 的张量，每行（100 行）正确的值（由索引中的相应索引选择？

小例子： 所以假设张量 A 是

[1, 2, 3]
[4, 5, 6]
[7, 8, 9]

张量 B 是

[0,2,1]

然后我想作为输出

[1,6,8]

Answer 1

您可以加入具有适当范围的 B 张量以创建二维索引（在您的示例中为 [[0, 0], [1, 2], [2, 1]]），然后使用 tf.gather_nd 提取元素：

b_2 = tf.expand_dims(b, 1)
range = tf.expand_dims(tf.range(tf.shape(b)[0]), 1)
ind = tf.concat(1, [range, b_2])
res = tf.gather_nd(a, ind)