在 Tensorflow 2.0 中用另一个张量索引张量的第 k 维

Question

我有一个张量probs，其形状为(None, None, 110)，代表LSTM 中的(batch_size, sequence_length, 110)。 我有另一个张量indices，其形状为(None, None)，其中包含要从probs 的第三维中选择的元素的索引。

我想使用indices 来索引张量probs。

Numpy 等价物：

k, j = np.meshgrid(np.arange(probs.shape[1]), np.arange(probs.shape[0]))
indexed_probs = probs[j, k, indices]

由于probs 中的shape[0] 和shape[1] 未知，因此tf.meshgrid() 不是一个选项。 我找到了tf.gather、tf.gather_nd 和tf.batch_gather，但它们似乎都没有按照我的意愿行事。

有人知道怎么做吗？

Answer 1

您可以像这样使用tf.gather_nd 做到这一点：

indexed_probs = tf.gather_nd(probs, tf.expand_dims(indices, axis=-1), batch_dims=2)

顺便说一句，在 NumPy 中你可以使用 np.take_along_axis 来做同样的事情：

indexed_probs = np.take_along_axis(probs, np.expand_dims(indices, axis=-1), axis=-1)[..., 0]