PyTorch 里做多轴索引，看起来就是方括号里多写几个维度的事，但实际一跑起来，各种报错能把人整懵。IndexError、形状对不上、维度莫名消失、复制了一份内存却浑然不觉......这些坑我基本都踩过一遍。这篇就把 PyTorch 多轴索引的常见写法和典型报错一次性捋清楚，帮你少走弯路。
1. 多轴索引全景：先搞清楚 PyTorch 里有哪些索引玩法
1.1 最基础的逗号分隔多轴索引
在处理张量的时候，很多人会下意识用tensor[0][1][2]这种“连连看”写法。这在低维小张量上偶尔能跑通，但在 PyTorch 里真正地道的做法是直接在一个方括号里用逗号隔开各个维度的索引：
import torch x = torch.arange(24).reshape(2, 3, 4) # x 的形状是 [2, 3, 4] print(x[1, 2, 3]) # tensor(23)，等价于 x[1][2][3] print(x[1, :, :]) # 取第 1 个“页面”的所有内容，等价于 x[1] print(x[:, 1, :]) # 取所有页面的第 1 行，形状为 [2, 4]逗号分隔的本质是把“每一维选什么”一次性表达清楚，代码可读性高很多，而且更好配合切片、步长和省略号。
1.2 三类进阶索引方式
多轴索引不只是写个整数，实际开发中更常用的是这种组合：
切片索引（Slice）：用start:end:step的形式，在某个维度上取连续或等间隔的一段。
print(x[:, :, ::2]) # 每个页面、每行、按步长 2 取列布尔掩码索引（Boolean Mask）：用一个与张量形状匹配的布尔张量，选择满足条件的元素。这个方法在数据筛选和条件过滤时极其好用。
mask = x > 10 print(x[mask]) # 输出一维张量，只包含大于 10 的那些元素花式索引（Fancy Indexing / 整数数组索引）：用一个整数张量或列表，指定要取的行列下标。它是实现“按指定顺序采样”的关键工具。
indices = torch.tensor([0, 2]) print(x[:, :, indices]) # 在最后一维上取第 0 列和第 2 列三种方式混用是实战常态，但也是报错高发区。后面的章节会逐个拆解为什么报错、怎么处理。
2. 深入理解多轴索引背后的机制：为什么有的结果降维，有的却保留维度
2.1 关键差异：整数索引会“降一维”，切片会“保维度”
这是新手最容易懵的地方。先看一个直观的例子：
x = torch.arange(12).reshape(3, 4) # 形状 [3, 4] a = x[:, 1] # 取所有行的第 1 列 # 得到形状 [3]，维度从 2 变成了 1 b = x[:, 1:2] # 取所有行的第 1 列，但保留“列”这个维度 # 得到形状 [3, 1]区别就在于：整数索引会把那一维“抽掉”，而切片（即使只切一个元素）也会保留一个长度可能为 1 的维度。这个机制经常导致你在写后续运算时莫名其妙地收到广播错误。
提示：如果想让代码逻辑稳定，建议尽量使用切片表达式写清楚保留维度的意图，比如
x[:, 1:2]而不是x[:, 1]，尤其在拼接和矩阵运算之前。
2.2 布尔掩码和多维整数索引的维度规则
布尔掩码索引的结果通常是“压平”的一维张量，这与 TensorFlow、NumPy 的行为一致。因为它本质上是在“筛选”，你无法预知选出来多少个元素，所以只能丢成一个一维向量。
而多维整数索引遵循另一套规则：索引数组的形状和结果形状有对应关系。比如：
rows = torch.tensor([[0, 1], [1, 2]]) cols = torch.tensor([[0, 1], [1, 2]]) result = x[rows, cols] # rows 和 cols 形状一致，都是一维或二维 # 结果形状就和索引数组的形状相同也就是说，高级索引结果的形状与你传入的索引张量的形状一致，而不是与被索引张量的形状一致。理解这点，才能准确预测输出形状。
2.3 视图（View）与副本（Copy）：别让性能问题悄悄出现
切片索引在多数情况下返回的是原张量的“视图”，共享底层内存。这意味着你修改切片结果，原张量也会跟着变。
x = torch.arange(12).reshape(3, 4) view_slice = x[:, 1:] view_slice[0, 0] = 999 print(x) # 你会发现 x[0, 1] 也变成了 999但布尔掩码索引和花式索引通常返回的是“副本”，修改它们不会影响原张量。这个差异在写数据预处理、做 Batch 采样时影响很大。
如果你不确定当前操作返回的是视图还是副本，可以检查输出张量的_base属性：
print(view_slice._base is not None) # True，说明是视图3. 多轴索引高频报错：原因、排查思路和解决代码
3.1IndexError: index out of range维度越界
现象：索引值超过了该维度长度，或者维度数量写多了。
x = torch.zeros(3, 4) # x[:, 5] -> IndexError # x[0, 0, 0] -> IndexError原因：维度数对不上，或者下标超出边界。这是多轴索引里出现频率最高的报错。
解决方法：先打印x.shape确认维度，再检查索引是否合法。我写代码时有个习惯：在高维张量上操作前，先用一行断言检查维度：
assert x.dim() == 2 and x.shape[-1] >= 6另外要特别留意“负数索引”的边界。PyTorch 支持负数索引，但如果你拿一个负数索引去切片或花式索引，它指向的是倒数第几位：
print(x[:, -1]) # 取最后一列，合法负数越界的典型场景是：索引值为 -N，同时该维度长度小于 N。这时 PyTorch 会直接报出IndexError，建议提前用clamp把索引约束到合法范围内：
indices = torch.clamp(raw_indices, 0, x.shape[-1] - 1)3.2 布尔掩码形状不匹配：mask shape ... doesn't match the broadcast shape
现象：
mask = torch.randint(0, 2, (3, 3)).bool() # 形状 [3, 3] x = torch.arange(12).reshape(3, 4) y = x[mask]这段代码会报错，因为mask的形状和x的形状不匹配。
原因：布尔掩码索引要求掩码和原张量维度一致（能广播也行），否则无法确定要筛哪些元素。
解决思路：确认掩码的形状与目标维度一致。如果你只想过滤某一个维度，可以先把掩码形状处理好，例如mask.unsqueeze(1)再加到多轴索引里。更好用的方式是记录“行下标”，然后用花式索引：
rows = torch.nonzero(x[:, 0] > 5).squeeze() result = x[rows] # 等价于按某列条件筛选所有行用nonzero将条件转化为行索引，既绕过了掩码形状匹配的烦恼，又能自然参与多轴索引，是实战中很实用的技巧。
3.3 “too many indices” 与“索引张量过多维”错误
现象：
y = torch.zeros(2, 3) idx = torch.tensor([[0, 1], [1, 0]]) y[idx, idx, idx]此时报错信息里面会出现维度相关的提示，核心问题是：你给一个 2 维张量传了 3 组索引，维度数量对不上。
原因：多轴索引的数量不能超过张量本身的维度数。如果确实需要对两个以上维度做花式索引，但张量维度不够，需要引入“伪批量维度”：
z = y.unsqueeze(0) # 变成 [1, 2, 3] # 然后就可以用三个维度的索引去操作3.4 高级索引组合限制：易踩坑的“部分高级索引”问题
PyTorch 文档中有一个很容易忽略的规则：当你混合使用切片和高级索引时，高级索引的维度会放到被切片维度之后。这会导致结果维度的顺序与直觉不一致。
x = torch.arange(24).reshape(2, 3, 4) # 混用：切片在第 0 维，花式索引在第 1、2 维 result = x[:, [0, 2], [1, 3]]这个写法有一定限制：如果要取多个点的组合，最好把两个索引分开写，不要都放在同一个方括号里，否则容易导致结果不是你想要的行列组合。推荐的做法：
result = x[:, [0, 2], :][:, :, [1, 3]]分两步完成多维采样，每一步都很直观，不易出错。
3.5 设备不一致报错：数据在 CPU、索引在 GPU
现象：
x = torch.randn(3, 4).cuda() idx = torch.tensor([0, 2]) # 默认在 CPU y = x[:, idx] # 报错：Expected all tensors to be on the same device原因：索引张量必须和被索引张量在同一设备上。
解决方法：直接把索引张量也移动到同一设备：
idx_gpu = idx.to(x.device) y = x[:, idx_gpu]或者尽可能使用tensor.to(device)统一设备环境。比如在数据加载阶段就把所有内容统一，也能省去大量处理步骤。
4. 实战拆解：多轴索引的典型应用场景
4.1 场景一：按条件从 Batch 中筛选样本
处理图像分类数据时，经常需要从一批样本里挑出某个类别的数据。
scores = torch.randn(8, 10) # 8 个样本，10 个类别分数 labels = torch.randint(0, 10, (8,)) # 每个样本的真实标签 # 提取所有标签为 3 的样本 target_mask = labels == 3 selected_scores = scores[target_mask]上面这个写法简洁有效。但更常见的情况是，你要筛选出标签为某个集合的所有样本：
target_classes = [3, 5, 7] class_mask = torch.isin(labels, torch.tensor(target_classes)) selected = scores[class_mask]torch.isin是批量多类别筛选的好帮手，远比多个条件叠加更清爽。
4.2 场景二：注意力矩阵中按位置取索引
Transformer 相关的代码里经常需要从特征序列中按指定位置收集向量。假设你的特征形状为[B, T, D]（批量、序列长度、特征维度），你需要取每个序列的第 2 和第 5 个位置的向量：
feat = torch.randn(4, 10, 64) positions = torch.tensor([1, 4]) # 第 2 个和第 5 个位置 # 先构造 batch 维的索引 batch_idx = torch.arange(feat.size(0)).unsqueeze(1) # [B, 1] batch_idx = batch_idx.expand(-1, positions.size(0)) # [B, 2] out = feat[batch_idx, positions.unsqueeze(0).expand(4, -1), :]这里的关键在于batch_idx和positions的形状一致，结果是[B, 2, D]，也就是每个序列对应位置的向量全部提出来了。
实际代码中，我通常会把这一步骤封装成函数：
def gather_indices(feat, positions): B, T, D = feat.shape batch_idx = torch.arange(B, device=feat.device).view(-1, 1).expand(-1, len(positions)) pos_idx = torch.as_tensor(positions, device=feat.device).view(1, -1).expand(B, -1) return feat[batch_idx, pos_idx, :]这样后续调用时，只需传入特征和位置列表，返回结果十分稳定。
4.3 场景三：更高效的替代方案
多轴索引确实灵活，但有些需求其实适合专用 API，性能更好、写法也更简练。
torch.gather(input, dim, index)：按索引沿着某个维度取值，适合批量收集；torch.index_select(input, dim, index)：按索引取整个子维，适合固定维度选择；torch.take(input, indices)：先压平再取索引，适合快速取一批位置的值。
例如，同样的“按位置取序列特征”需求，用gather可以这样写：
indices = positions.unsqueeze(0).unsqueeze(-1).expand(4, -1, 64) # [B, K, D] out = torch.gather(feat, dim=1, index=indices)对比之下，gather对“同形索引张量”的构造要求比较多，但一旦写熟，性能比多层索引更稳定。
index_select则适合整维度选取：
out = torch.index_select(feat, dim=1, index=positions)它只作用于单一维度，不展开索引之间的组合关系，因此语义更简单。
5. 报错排查速查表与避坑清单
| 报错类型 | 常见信息 | 直接原因 | 处理方案 |
|---|---|---|---|
| 维度越界 | IndexError: index out of range | 下标超出该维长度 | 检查 shape，或对索引做clamp |
| 维度数量过多 | too many indices | 索引维度数量超过张量维度 | 打印 dim，用unsqueeze补维度 |
| 掩码形状不匹配 | mask shape ... doesn't match | 布尔掩码与原张量形状不一致 | 利用nonzero转行索引，或用unsqueeze对齐 |
| 设备不一致 | Expected all tensors to be on the same device | 索引张量设备与原张量不同 | idx.to(x.device) |
| 高级索引组合异常 | 结果形状不像预期的那样 | 混合切片和高级索引导致维度顺序变动 | 拆分索引步骤，明确每一步结果 |
| 返回视图/副本混淆 | 修改后原张量也变了 | 切片索引返回视图 | 显式使用.clone()保证隔离 |
以下经验是我反复踩坑后总结出来的，想省时间可以直接记下来：
- 先看 shape，再写索引。多轴索引的一切问题，归根结底都是形状问题。代码里多打印两次
shape不丢人。 - 保持设备统一。每次创建索引张量，我都直接写成
torch.tensor(..., device=x.device)，避免后面再补.to()。 - 能用专用 API 就尽量用专用 API。
index_select、gather等函数经过优化，在很多场景下比手工构造索引矩阵更省心。 - 代码里明确区分视图和副本。做数据预处理或需要原地修改时，避免对视图操作产生副作用，一旦不确定就
.clone()。
PyTorch 多轴索引在大部分时候的“报错”并不是框架有 bug，而是我们没把它隐含的“降维”和“维度顺序”规则摸透。把这些规则理清楚，再配合这上面几个实战技巧，你在遇到这类问题时基本能秒定位、快速改。
