本文深入讲解如何在 PyTorch 中利用形状为 [b, k] 的索引张量 B，对形状为 [b, m, n] 的高维张量 A 执行高效批量索引，最终得到 [b, k, n] 的输出。核心思路在于合理扩展索引维度并配合 torch.gather 实现精准的逐行抽取。
很多人处理高维张量的批量索引时都会遇到瓶颈——尤其是当索引本身是 batch 维独立的二维张量，而目标张量还需要保留最后一维的所有列。直接使用 torch.index_select 或 torch.take 会发现根本行不通：它们只接受一维索引。而 torch.gather 虽然能处理多维输入，但要求输入张量与索引在除指定维度外的所有维度上严格对齐。那么，如何用形状为 [b, k] 的索引 B，从形状为 [b, m, n] 的张量 A 中提取每 batch 独立的 k 行，同时完整保留最后一维 n 的全部数值？
关键在于对索引张量进行升维并使其与目标张量形状对齐。具体步骤拆解如下：
-
明确索引语义——我们期望的输出结果满足
out[b, k, n] == A[b, B[b, k], n]，即对每个 batch b，从A[b]（形状[m, n]）中挑选出B[b, k]所指定的第 k 行，总共取出 k 行，且每行保留完整的 n 列信息。 -
扩展索引维度：先将
B（形状[b, k]）转换为[b, k, 1]，再通过广播机制扩展为[b, k, n]，从而与A的最后一维匹配：B_expanded = B.unsqueeze(-1).expand(-1, -1, A.size(-1)) # [b, k, n]
-
传入 torch.gather：沿着
dim=1（即 m 维度）执行 gather 操作。此时A的形状为[b, m, n]，B_expanded的形状为[b, k, n]，除 gather 维度外其他维度均已对齐：out = torch.gather(A, dim=1, index=B_expanded) # 输出 shape: [b, k, n]
下面提供一个完整可运行的示例，方便直接验证实现效果：
import torch
b, m, n, k = 2, 5, 4, 3
A = torch.randn(b, m, n) # [2, 5, 4]
B = torch.randint(0, m, (b, k)) # [2, 3]，值 ∈ [0, 4]
# 扩展索引：[b,k] → [b,k,1] → [b,k,n]
B_idx = B.unsqueeze(-1).expand(-1, -1, n)
# 沿 dim=1 gather
out = torch.gather(A, dim=1, index=B_idx)
print(f"A.shape: {A.shape}") # torch.Size([2, 5, 4])
print(f"B.shape: {B.shape}") # torch.Size([2, 3])
print(f"out.shape: {out.shape}") # torch.Size([2, 3, 4])
# 验证：out[0,0] 应等于 A[0, B[0,0]]
assert torch.equal(out[0, 0], A[0, B[0, 0]])
有几个细节值得特别留意：
- 索引张量
B中的每个数值必须严格落在[0, m)区间内，否则会触发IndexError； torch.gather不支持负索引（与 NumPy 不同），使用前需确保索引为非负；- 该操作完全可微，若需要回传梯度可直接使用。如果仅用于推理阶段并希望加速，也可考虑
torch.nn.functional.embedding——只需将A视为 embedding 权重，将B视为 token IDs 即可； - 当然也能用循环配合
torch.index_select实现，但那样会失去向量化优势，性能相差悬殊，通常不推荐。
掌握这一模式后，像序列抽取、top-k 特征筛选、动态掩码选择等常见场景，都能轻松应对。
