如何从TensorFlow张量的每行提取指定元素？

给定TensorFlow张量：

import tensorflow as tf

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

以及索引张量：

i = tf.convert_to_tensor([[0],[2]])

需要从t中提取对应元素，得到结果[[1],[7]]。你当前使用拼接索引对的方法可行，这里提供两种更简洁高效的解法：

解法一：直接张量索引 #

通过生成行索引、压缩列索引维度后，直接利用张量的索引语法提取元素，再调整维度匹配目标形状：

# 生成对应每一行的索引 [0, 1]
row_indices = tf.range(tf.shape(t)[0])
# 将列索引从 (2,1) 压缩为 (2,) 的 [0, 2]
col_indices = tf.squeeze(i, axis=1)
# 提取元素并增加最后一维，得到 (2,1) 的结果张量
result = t[row_indices, col_indices][..., tf.newaxis]
print(result.numpy())  # 输出：[[1] [7]]

解法二：使用tf.batch_gather #

tf.batch_gather专为批量场景下的索引提取设计，可直接传入原张量和列索引张量，一步得到目标结果：

result = tf.batch_gather(t, i)
print(result.numpy())  # 输出：[[1] [7]]

这两种方法都无需手动拼接索引对，代码更简洁直观，且tf.batch_gather是TensorFlow针对此类场景优化的API，性能表现更稳定。

内容的提问来源于stack exchange，提问作者Alberto