您可以使用NumPy的take_along_axis的TensorFlow变体，

tf.experimental.numpy.take_along_axis(t, i, axis=1)

您可以简单地将i与tf.range(...)堆叠，如下所示

import tensorflow as tf

t = tf.convert_to_tensor([
  [1,2,3,4],
  [5,6,7,8]
])
i = tf.convert_to_tensor([0, 2])

length = tf.shape(i)[0]
indices = tf.stack([tf.range(length), i], axis=1)
# [0, 0], [1, 2]]

tf.gather_nd(t, indices)
# [1, 7]

我不确定有没有更好的解决办法。