python - 如何在 tensorflow 中收集带有索引的元素

Question

例如，

import tensorflow as tf

index = tf.constant([[1],[1]])
values = tf.constant([[0.2, 0.8],[0.4, 0.6]])

如果我使用extract = tf.gather_nd(values, index)返回的是

[[0.4 0.6]
 [0.4 0.6]]

但是，我想要的结果是

[[0.8], [0.6]]

其中索引沿 axis = 1，但是 tf.gather_nd 中没有轴参数设置。

我该怎么办？谢谢!

Answer 1

将范围连接到索引:

index = tf.stack([tf.range(index.shape[0])[:, None], index], axis=2)
result = tf.gather_nd(values, index)

result.eval(session=tf.Session())
array([[0.8],
       [0.6]], dtype=float32)