Python 矩阵张量索引

我有矩阵（3 x 15）

假人

，其中标记序列为行：

[[ 1 66 67 68  0  0  0  0  0  0  0  0  0  0  0]
[ 1 66 67 66 68 66 67 66  0  0  0  0  0  0  0]
[ 1 66 67 68 18 19 20 21 22 23 24 25 26 17  0]]

此外，还有一个形状（3 x 15 x n_标记）的张量

probs

，具有标记概率

从

probs

中，我只需要选择

假人中令牌的概率

我认为，可以使用矩阵作为张量的索引，但我还没有找到如何做到这一点。
你可以这样做：

import tensorflow as tf

dummies = ...
probs = ...
s = tf.shape(dummies)
i = tf.range(s[0])
j = tf.range(s[1])
ii, jj = tf.meshgrid(i, j, indexing='ij')
idx = tf.stack([ii, jj, dummies], axis=-1)
result = tf.gather_nd(probs, idx)

我在
结果中有5个维度
@koPytok抱歉，修复了它，它应该是
tf.gather\u nd
，而不是
tf.gather
。