如何在张量流中收集具有索引的元素

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) return 是

[[0.4 0.6]
 [0.4 0.6]]

不过，我想要的结果是

[[0.8], [0.6]]

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

我该怎么办？谢谢！

Answer 1

将范围连接到 index:

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)

how to gather element with index in tensorflow