TensorFlow：利用argmax索引切片同形状张量

嘿，这个需求很常见，用TensorFlow的内置函数就能轻松实现，我给你两种实用的方法，你可以根据习惯选择：

方法一：使用tf.batch_gather（简洁高效） #

这种方法最适配你这种在最后一维取索引的场景，步骤非常直观：

1. 先把max_ind扩展一个维度，让它的形状从[1, 1000]变成[1, 1000, 1]，确保和t2的维度匹配；
2. 用tf.batch_gather根据扩展后的索引从t2中提取对应元素；
3. 最后去掉多余的维度，得到目标形状[1, 1000]。

代码示例：

# 假设你已经得到了max_ind = tf.argmax(t1, axis=-1)
max_ind_expanded = tf.expand_dims(max_ind, axis=-1)  # 形状变为[1, 1000, 1]
result = tf.batch_gather(t2, max_ind_expanded)       # 形状[1, 1000, 1]
result = tf.squeeze(result, axis=-1)                 # 最终形状[1, 1000]

方法二：使用tf.gather_nd（通用灵活） #

如果之后需要处理更复杂的多维度索引场景，tf.gather_nd会是更通用的选择。它需要我们构造每个元素的完整坐标（即batch、seq、feat三个维度的索引），再根据坐标提取元素：

代码示例：

# 获取各维度的大小
batch_size, seq_len, _ = tf.shape(t1)[0], tf.shape(t1)[1], tf.shape(t1)[2]

# 生成batch维度的索引：形状[1, 1000]
batch_idx = tf.tile(tf.expand_dims(tf.range(batch_size), 1), [1, seq_len])
# 生成seq维度的索引：形状[1, 1000]
seq_idx = tf.tile(tf.expand_dims(tf.range(seq_len), 0), [batch_size, 1])

# 拼接成完整的坐标：形状[1, 1000, 3]，每个元素是[batch_idx, seq_idx, feat_idx]
indices = tf.stack([batch_idx, seq_idx, max_ind], axis=-1)

# 根据坐标提取t2中的元素，得到目标形状[1, 1000]
result = tf.gather_nd(t2, indices)

验证结果 #

你可以用一个小张量测试正确性，比如：

t1 = tf.constant([[[1, 3, 2], [0, -1, 5]]])  # 形状[1,2,3]
max_ind = tf.argmax(t1, axis=-1)             # 结果是[[1, 2]]
t2 = tf.constant([[[10, 20, 30], [40, 50, 60]]])
# 用上面的方法得到的结果应该是[[20, 60]]，和t1最大值位置对应的t2元素完全一致

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