如何使用tf.gather按指定轴的对应单个索引提取元素？

你当前的代码是沿着axis=1（列方向）重复提取指定列，所以每一行都会取indices里所有索引对应的列，最终得到3×3的结果。而你需要的是逐行匹配索引提取单个元素——第i行取indices[i]位置的元素，最终得到一维数组。

这里有两种直观的解决方案：

方法一：使用tf.gather_nd（推荐） #

tf.gather_nd支持传入多维索引，刚好适合这种“行号+列号”的配对提取场景：

import tensorflow as tf
indices = tf.constant([0, 1, 1])
x = tf.constant([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

# 生成与indices长度一致的行索引：[0, 1, 2]
row_indices = tf.range(tf.shape(indices)[0])
# 将行索引与列索引配对成二维索引：[[0,0], [1,1], [2,1]]
gather_indices = tf.stack([row_indices, indices], axis=1)
# 提取对应位置的元素
result = tf.gather_nd(x, gather_indices)

with tf.Session() as sess:
    selection = sess.run(result)
    print(selection)

运行结果：

[1 5 8]

方法二：转换为一维索引后用tf.gather #

如果一定要用tf.gather，可以先把多维数组展平，计算每个目标元素的一维索引后再提取：

import tensorflow as tf
indices = tf.constant([0, 1, 1])
x = tf.constant([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

row_indices = tf.range(tf.shape(indices)[0])
# 计算一维索引：行号 × 列数 + 列索引
flat_indices = row_indices * tf.shape(x)[1] + indices
# 展平数组后提取元素
result = tf.gather(tf.reshape(x, [-1]), flat_indices)

with tf.Session() as sess:
    selection = sess.run(result)
    print(selection)

这个方法也能得到相同的结果，但tf.gather_nd的写法更贴合多维数组的操作逻辑，可读性更强。

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