如何在TensorFlow中用(n-k)维张量索引n维张量？

不管你用PyTorch还是TensorFlow，都可以通过专门的索引函数或者构造完整索引的方式实现需求，以下是具体实现：

PyTorch 实现 #

方法1：使用torch.gather（推荐） #

gather是PyTorch中专门用于按指定维度索引收集元素的函数，需要先将索引张量B的形状扩展到与A的维度匹配：

import torch

# 示例张量
A = torch.randn(7, 7, 2, 4)
B = torch.argmax(A, dim=2)  # 形状[7,7]，取值0或1

# 扩展B的维度，使其在dim=2以外的维度与A匹配
B_expanded = B.unsqueeze(-1).unsqueeze(-1)  # 变为[7,7,1,1]
# 按dim=2收集元素，之后可以squeeze去掉多余维度
C = A.gather(dim=2, index=B_expanded).squeeze(2)  # 结果形状[7,7,4]
# 如果需要[7,7,1,4]的形状，去掉squeeze即可：
# C = A.gather(dim=2, index=B_expanded)

方法2：构造完整索引（直观） #

手动生成每个位置的第0、1维索引，结合B的第2维索引，直接提取对应元素：

# 生成i,j网格索引，形状均为[7,7]
i = torch.arange(7).unsqueeze(1).repeat(1, 7)
j = torch.arange(7).unsqueeze(0).repeat(7, 1)

# 按索引提取，结果形状[7,7,4]
C = A[i, j, B, :]

TensorFlow 实现 #

方法1：使用tf.gather_nd #

构造包含所有维度索引的数组，直接提取目标元素：

import tensorflow as tf

# 示例张量
A = tf.random.normal((7, 7, 2, 4))
B = tf.argmax(A, axis=2)  # 形状[7,7]，取值0或1

# 生成i,j网格索引
i, j = tf.meshgrid(tf.range(7), tf.range(7), indexing='ij')
# 组合成完整索引：每个位置为[i,j,B[i,j]]
indices = tf.stack([i, j, B], axis=-1)
# 提取元素，结果形状[7,7,4]
C = tf.gather_nd(A, indices)
# 如果需要[7,7,1,4]的形状，扩展维度即可：
# C = tf.expand_dims(C, axis=2)

方法2：使用tf.gather #

通过扩展B的维度并指定batch维度，实现按索引收集：

# 扩展B的维度至[7,7,1,1]
B_expanded = tf.expand_dims(tf.expand_dims(B, axis=-1), axis=-1)
# 按axis=2收集，指定前2个维度为batch维度无需索引
C = tf.gather(A, B_expanded, axis=2, batch_dims=2)  # 结果形状[7,7,1,4]
# 若要[7,7,4]，可squeeze：
# C = tf.squeeze(C, axis=2)

为什么直接用A[B]无效？ #

A[B]会将B中的每个元素当作第0维的索引来取值，也就是从A的第0维（形状为7）中选取对应下标的元素，这完全不符合你要按第2维索引的需求，因此必须明确指定索引对应的维度，使用上述专门的索引方法。

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