python - 使用 Tensorflow 2.0 中的另一个张量索引张量的第 k 维

Question

我有一个张量 probs具有形状 (None, None, 110)代表(batch_size, sequence_length, 110)在 LSTM 中。
我有另一个张量 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]

自 shape[0]和 shape[1]的 probs未知， 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]