张量索引这件事，说简单也简单——不就是拿中括号取数据嘛；说复杂也复杂，我见过不少写了半年PyTorch的人，遇到x[:, None]、x[..., 0]、x[mask]这类写法还是得停下来想半天，更别提index_select、gather、scatter这几个函数了，每次用都得翻文档。索引操作是张量操作里出现频率最高的动作之一，数据切片、维度对齐、标签映射、注意力掩码，几乎每一处都离不开它。这篇就把PyTorch张量索引这一块从头到尾捋一遍，从最基础的中括号取值，到高级索引的广播规则，再到几个容易混淆的索引函数之间的区别，配上能直接跑的实验代码和我在实际项目里踩过的坑。不管你是刚接触PyTorch想打牢基础，还是已经能写模型但索引写得不够利索，下面这些内容应该都能对上你的需求。
1. 先搞清楚PyTorch索引到底在索引什么
很多人学索引的第一反应是背语法，x[0]取第一行、x[:, 1]取第二列，背完就上手写代码。这样做的后果是，一旦维度超过三维、或者索引里混进了切片和整数，脑子里的"行列表格"图像就崩了。所以这一节先把索引的底层逻辑讲透，后面所有花式写法都是从这条逻辑推出来的。
1.1 张量的本质是"带步长的内存视图"
PyTorch的张量并不是一个二维表格或者三维立方体，它在内存里就是一维连续排列的一串数字，所谓"形状"只是给这串数字套了一层坐标解释。每个维度对应一个步长（stride），表示沿着这个维度走一步，在内存里要跳过多少个元素。举个具体例子：
import torch x = torch.arange(24).reshape(2, 3, 4) print(x.shape) # torch.Size([2, 3, 4]) print(x.stride()) # (12, 4, 1)这个(12, 4, 1)的意思是：第0维走一步跳过12个元素，第1维走一步跳过4个，第2维走一步跳过1个。你写x[1, 2, 3]，PyTorch实际算的是内存偏移1*12 + 2*4 + 3*1 = 23，直接取第23号元素。理解了这一点，你就能明白为什么索引操作几乎不复制数据——它只是换了一套"坐标到偏移"的映射规则，返回的是原内存的新视图。
提示：视图意味着共享内存。对切片结果做原地修改（如
y.add_(1)）会改到原张量，这是新手最常见的"数据莫名其妙变了"的来源。需要独立副本时用.clone()。
1.2 基本索引：整数、切片、省略号三件套
基本索引就是Python列表那套东西的延伸，但多了两个PyTorch特有的符号：...（省略号）和None（新增维度）。整数索引会降维，切片索引保持维度，这是两者最核心的区别，也是后面很多"形状对不上"报错的根源。
x = torch.arange(24).reshape(2, 3, 4) print(x[0].shape) # torch.Size([3, 4]) 整数索引，第0维被吃掉 print(x[0:1].shape) # torch.Size([1, 3, 4]) 切片索引，维度保留 print(x[..., 0].shape) # torch.Size([2, 3]) 省略号代表前面所有维度 print(x[:, None, :, :].shape) # torch.Size([2, 1, 3, 4]) None插入新维度None这个用法值得单独说。它等价于unsqueeze，在需要手动对齐广播维度时极其好用。比如你有一个形状(B,)的标签张量，想和形状(B, C)的预测做逐样本运算，写成labels[:, None]变成(B, 1)就能自动广播，比先unsqueeze(1)再运算少敲好几个字符。我个人的习惯是：临时对齐维度用None，需要保留语义清晰的代码用unsqueeze，两者性能完全一样。
...的价值在于"我不想数前面有几个维度"。写通用函数时尤其有用，比如一个处理任意维特征图的归一化函数，用x[..., :3]取前三个通道，不管输入是三维还是四维都不会写错。反过来，如果你硬写x[:, :3]，输入变成四维时就取错位置了。
1.3 负索引与越界：边界规则要记牢
负索引从末尾倒数，x[-1]是最后一个，这个和Python一致。但PyTorch在越界处理上和列表有个重要差异需要留意：基本索引越界会直接报错，不会像某些库那样静默截断。
x = torch.arange(10) print(x[-3:]) # tensor([7, 8, 9]) 负索引切片正常 print(x[5:100]) # tensor([5, 6, 7, 8, 9]) 切片越界自动截断，不报错 # print(x[100]) # IndexError: index 100 is out of bounds切片越界截断、整数索引越界报错，这个不对称行为要记住。实际写数据加载逻辑时，如果索引来自外部输入（比如文件里读出来的样本编号），建议先做范围校验或者用torch.clamp夹一下，否则训练跑了几小时突然崩在一个越界索引上，排查起来很烦。
2. 高级索引：花式取值的规则与广播陷阱
基本索引只能取"连续的一块"或者"固定位置的单个元素"，一旦你需要按任意顺序取一批元素、按条件筛选、或者按索引表查表，就得上高级索引。高级索引是PyTorch索引里最容易出错的部分，核心难点在于多个索引张量之间的广播规则，这一节把它拆开讲。
2.1 整数数组索引：成对取值的直觉
用整数张量做索引，语义是"按给定坐标逐个取元素"。当只有一个索引张量时很直观：
x = torch.arange(12).reshape(3, 4) idx = torch.tensor([0, 2, 1]) print(x[idx]) # 取第0、2、1行，形状 (3, 4) print(x[idx].shape) # torch.Size([3, 4])注意这里返回的形状是(3, 4)——索引张量的形状(3,)替换掉了被索引的第0维，其余维度原样保留。这是高级索引的一条基本规律：索引张量的形状会"顶替"被索引维度的位置。
当有多个索引张量时，规则变成"按位置配对"：
rows = torch.tensor([0, 1, 2]) cols = torch.tensor([1, 2, 3]) print(x[rows, cols]) # tensor([1, 6, 11]) 取(0,1)、(1,2)、(2,3)这里取的是三个对角位置，而不是三行三列组成的子矩阵。想要子矩阵得用x[rows][:, cols]或者x[rows[:, None], cols]。这个区别是新手第一个大坑：x[a, b]是配对取值，x[a][:, b]才是笛卡尔积式的行列交叉。我见过有人想取一个batch里每个样本对应标签位置的logit，写成logits[batch_idx, label_idx]——这个写法恰好是对的，因为就是要配对；但如果他想取所有样本对所有标签的完整子矩阵，这个写法就错了。
2.2 索引张量的广播：形状不一致时会发生什么
多个索引张量不要求形状完全相同，PyTorch会先对它们做广播，再配对取值。这条规则威力很大，但也容易写出"看起来对、结果形状很怪"的代码：
x = torch.arange(12).reshape(3, 4) rows = torch.tensor([[0], [1]]) # 形状 (2, 1) cols = torch.tensor([0, 2, 3]) # 形状 (3,) print(x[rows, cols].shape) # torch.Size([2, 3])(2,1)和(3,)广播成(2,3)，于是取出一个2x3的结果，第i行第j列是x[rows[i], cols[j]]。这个机制在实现"每个样本取多个候选位置"时非常顺手，比如从形状(B, N)的打分矩阵里，为每个样本取top-k位置的值，可以构造(B, 1)的样本索引和(B, k)的位置索引，一次索引搞定，不用写循环。
但广播也有反直觉的地方：如果两个索引张量形状既不相同也不能广播，PyTorch会直接报错，报错信息里会提示"shape mismatch"。遇到这种报错，先检查两个索引张量的形状，用None手动补维度往往能解决。
2.3 布尔掩码索引：条件筛选的标准姿势
布尔索引是数据清洗和损失计算里用得最多的高级索引。传入一个和某维度等长的布尔张量，返回所有为True位置对应的元素：
x = torch.randn(3, 4) mask = x > 0 print(mask.shape) # torch.Size([3, 4]) print(x[mask].shape) # 一维，长度等于True的个数关键点：布尔索引返回的永远是一维张量，不管原张量几维、掩码几维。因为它是"把所有满足条件的位置摊平收集起来"，位置信息在结果里丢失了。如果你需要保留形状，得用torch.where或者masked_fill：
# 保留形状，不满足条件的位置填0 y = torch.where(mask, x, torch.zeros_like(x)) # 或者原地填充 z = x.clone() z.masked_fill_(~mask, 0.0)布尔掩码还有一个高频用法是"按条件给损失加权"，比如只对非padding位置计算loss：
loss = criterion(pred, target) # 逐元素loss，形状 (B, L) loss = loss[target != pad_token_id] # 只保留有效位置 final_loss = loss.mean()这里target != pad_token_id生成布尔掩码，索引后摊平成一维，再求均值。这个写法比先乘掩码再除以有效个数要简洁，而且不会有"除以零"的风险（掩码全False时loss为空，mean()返回nan，需要额外判断，这是个小坑）。
注意：布尔掩码索引的结果是副本，不是视图。对
x[mask]做原地修改不会影响x，这点和切片不同，别搞混。
2.4 混合索引：基本索引和高级索引混用时的维度顺序
真实代码里很少纯用高级索引，更多是"前面几维切片、后面几维用索引张量"的混合写法。混合索引的规则是：高级索引部分的结果维度会被放到前面，基本索引部分保持相对顺序。这条规则不看文档基本猜不到：
x = torch.arange(24).reshape(2, 3, 4) idx = torch.tensor([0, 2]) print(x[:, idx].shape) # torch.Size([2, 2, 4]) 切片在前，索引替换第1维 print(x[idx, :, 1].shape) # torch.Size([2, 4]) 高级索引被提到最前第二个例子里，idx索引第0维、:切片第1维、1是整数索引第2维。按"高级索引提前"的规则，结果是(2, 4)——索引张量的2在前，切片保留的3被整数索引吃掉变成4。如果你预期的是(3, 2)，就会对不上。我的建议是：混合索引超过两个维度时，拆成两步写，先切片再索引，形状一目了然，可读性远好于一行塞满各种符号。
3. 那几个长得像的索引函数，到底该用哪个
index_select、gather、take、masked_select、scatter——这几个函数名字相近、功能有重叠，很多人是"哪个能跑通用哪个"，结果代码又慢又难读。这一节按"解决什么问题"来区分它们，给出选型判断标准。
3.1 index_select 与 take：按索引表取值的两种粒度
index_select沿指定维度按索引表取值，语义和x[idx]在单维度上完全一致，但它是函数形式，可以明确指定维度：
x = torch.arange(12).reshape(3, 4) idx = torch.tensor([0, 2]) print(torch.index_select(x, dim=0, index=idx).shape) # (2, 4) print(torch.index_select(x, dim=1, index=idx).shape) # (3, 2)take则是把张量摊平后按索引取，等价于x.flatten()[idx]：
print(torch.take(x, torch.tensor([0, 5, 11]))) # tensor([0, 5, 11])选型上很简单：需要指定维度就用index_select，需要跨维度按扁平位置取就用take。take在实现自定义采样、按全局位置取特征时有用，但日常用得少。index_select在embedding查表、按类别取权重时很常见，而且它在某些后端上有专门的优化，比x[idx]在特定形状下略快，不过差异通常可以忽略，选它更多是为了代码语义清晰。
3.2 gather：按"每个位置各自的索引"取值
gather是这几个函数里最容易被误解的。它的语义不是"取一批行"，而是"输出张量的每个位置，都去输入张量里按对应索引取一个值"。索引张量和输出同形状，每个元素指定沿某维度取哪个位置：
x = torch.tensor([[10, 20, 30], [40, 50, 60]]) idx = torch.tensor([[2, 1, 0], [0, 2, 1]]) print(torch.gather(x, dim=1, index=idx)) # tensor([[30, 20, 10], # [40, 60, 50]])第0行按[2,1,0]重排，第1行按[0,2,1]重排。gather的典型场景是从概率分布里取目标类别的概率：probs.gather(1, labels[:, None])取出每个样本对应标签的概率，形状(B, 1)。这个用法在强化学习、序列生成里到处都是。
用gather有两个硬性约束必须记住：索引张量和输入张量维度数相同，且除了dim指定的维度外，其余维度大小必须一致（或者能广播）。违反任何一条都会报错。另外索引值必须在合法范围内，gather不做越界检查，越界会读到未定义内存，这是它和基本索引的一个重要区别——gather越界不报错，结果不可预测，索引来自外部数据时务必先校验。
3.3 scatter 与 masked_select：写入与筛选的对应操作
scatter是gather的逆操作：把源张量的值按索引"撒"到目标张量的对应位置。
target = torch.zeros(2, 3) src = torch.tensor([[1, 2, 3], [4, 5, 6]]) idx = torch.tensor([[0, 1, 2], [2, 1, 0]]) target.scatter_(1, idx, src) # tensor([[1., 2., 3.], # [6., 5., 4.]])scatter在实现one-hot、按索引累加统计量、把稀疏更新写回稠密张量时非常有用。有个高频变体scatter_add_，相同索引位置的值会累加而不是覆盖，做直方图统计、按类别汇总损失时用它一行搞定。
masked_select就是布尔索引的函数形式，x[mask]和torch.masked_select(x, mask)结果一致。函数形式的好处是可以显式传out参数复用内存，在循环里频繁筛选时能减少分配开销。
| 函数 | 语义 | 输出形状 | 典型场景 |
|---|---|---|---|
index_select | 沿指定维度按索引表取 | 索引维被替换 | embedding查表、按类取权重 |
take | 摊平后按位置取 | 一维 | 全局位置采样 |
gather | 每位置按各自索引取 | 与索引张量同形 | 取目标类别概率、重排 |
scatter | 按索引写入 | 与目标同形 | one-hot、稀疏更新回写 |
masked_select | 布尔筛选 | 一维 | 条件筛选、有效位置收集 |
4. 索引在真实模型代码里的高频用法
光看语法容易忘，索引这东西必须放到具体场景里才有记忆点。这一节挑几个模型代码里反复出现的索引模式，讲清楚它们为什么这么写、有没有更优写法。
4.1 注意力掩码：布尔索引与where的配合
自注意力里对padding位置做mask是索引的经典应用。假设attn形状(B, L, L)，mask形状(B, L)标记有效位置，标准写法是：
# 把无效位置的注意力分数置为负无穷，softmax后自然为0 attn = attn.masked_fill(~mask[:, None, :], float('-inf')) attn = attn.softmax(dim=-1)这里mask[:, None, :]把(B, L)扩成(B, 1, L)，靠广播对齐到(B, L, L)。用masked_fill而不是布尔索引，是因为要保留形状。如果写成attn[~mask] = -inf，形状对不上会直接报错——布尔掩码必须和整个张量同形，或者只索引某一维。这个细节很多人第一次写注意力时都会卡一下。
另一个常见需求是"按mask把变长序列收集成扁平张量"，比如计算loss前把padding去掉：
flat = hidden[mask] # (有效token总数, D)这个写法比循环每个样本、按长度切片要快得多，而且天然支持batch内变长。代价是丢失了样本边界信息，如果后续需要按样本还原，得配合cumsum记录每段长度。
4.2 标签映射与embedding查表
embedding层本质就是一次index_select：weight[idx]取出idx对应的行。理解这一点后，很多"查表"需求都可以直接用索引实现，不必非得建一个nn.Embedding。比如有一个形状(num_classes, D)的类别原型张量，想按每个样本的类别取原型：
protos = torch.randn(10, 64) # 10个类别原型 labels = torch.tensor([3, 7, 1, 3]) # 每个样本的类别 selected = protos[labels] # (4, 64)这比torch.index_select(protos, 0, labels)更简洁，性能一样。需要注意的是labels必须是long类型（torch.int64），用int32会报"expected scalar type Long"的错，这是从numpy转过来时的高频问题——numpy默认int64还好，但从某些数据管道出来的可能是int32，记得.long()一下。
4.3 按索引聚合：scatter_add做分段求和
变长序列的池化、按用户ID汇总特征这类"分段求和"需求，用scatter_add_可以完全避免循环：
# 每个token属于某个样本，把token特征按样本求和 feat = torch.randn(100, 32) # 100个token的特征 owner = torch.tensor([...]) # 长度100，每个token属于哪个样本(0~B-1) B = owner.max().item() + 1 pooled = torch.zeros(B, 32) pooled.scatter_add_(0, owner[:, None].expand(-1, 32), feat)owner[:, None].expand(-1, 32)把(100,)扩成(100, 32)，让每个token的32维特征都加到对应样本行。这个模式在实现"按段落聚合句子向量"、"按会话聚合消息特征"时非常实用，比for循环加index_add快一个数量级。index_add_是另一个选择，语义是"按索引把源张量加到目标"，写法上pooled.index_add_(0, owner, feat)更简洁，效果等价，我一般优先用index_add_，可读性更好。
5. 踩过的坑与性能注意事项
索引写错往往不报错，只是结果悄悄不对，这类bug最难查。这一节列几个我实际遇到过、也在别人代码里见过多次的问题，附上排查思路。
5.1 视图与副本混淆导致的数据污染
前面提过切片是视图，但很多人低估了它的影响范围。一个真实场景：从一个大batch里切出子batch做某种变换，以为改的是子batch，结果原数据被改了。
sub = big_batch[0:4] # 视图 sub.mul_(2) # 原地操作，big_batch前4行也被乘2排查这类问题的办法是养成习惯：凡是切片结果要修改，先.clone()。反过来，如果只是读取，视图反而省内存，不用多此一举。判断标准很简单——后面有没有原地操作（带下划线的方法、+=、赋值到切片）。有就clone，没有就不clone。
另一个隐蔽的点是reshape和view的区别。view要求内存连续，不连续时会报错；reshape在必要时会复制。如果你对一个转置过的张量调view，报错信息是"view size is not compatible with input tensor's size and stride"，这时候要么先.contiguous()，要么直接用reshape。我一般默认用reshape，省心。
5.2 gather/scatter的越界与形状陷阱
gather越界不报错这一点前面提过，实际影响比想象中大。如果索引来自模型输出（比如argmax结果），理论上不会越界；但如果索引来自外部标注、或者经过某种变换后没做clamp，就可能读到垃圾值，表现为loss突然变成nan或者异常大。排查时可以在gather前加一句断言：
assert idx.min() >= 0 and idx.max() < x.size(dim)开发阶段加上，上线前去掉，能省很多调试时间。scatter同理，索引越界会静默丢弃或写到错误位置。
形状陷阱方面，gather要求索引张量和输入维度数相同，很多人从(B,)的索引直接传给gather会报"Index tensor must have the same number of dimensions as input tensor"，解决方法是idx.unsqueeze(-1)补一维。这个报错信息其实挺明确，但第一次遇到容易懵。
5.3 高级索引的性能：什么时候该换成函数形式
高级索引x[idx]在底层会走一套通用的索引内核，某些形状下比专用函数慢。实测下来，差异主要体现在两个场景：一是索引张量很大且需要跨步访问时，index_select通常更快；二是需要重复用同一个索引时，把索引预先转成long并保证连续，能减少每次调用的转换开销。
不过说实话，绝大多数模型代码里索引不是瓶颈，没必要为了这点差异把可读性牺牲掉。真正需要注意的是在训练循环里反复创建大索引张量，比如每个step都torch.tensor([...])构造索引，这个构造开销可能比索引本身还大。解决办法是把索引张量在循环外建好，循环内复用。
还有一个容易忽略的点：布尔索引x[mask]在mask很稀疏（True很少）时，结果很小但中间过程要扫描整个张量；如果mask极度稀疏且你只需要少数几个位置，用mask.nonzero()拿到位置再索引，有时更快。这个优化比较极端，一般用不上，知道有这么个思路就行。
6. 把索引规则内化成直觉的练习路径
索引这东西，看十遍不如写一遍。最后分享一套我自己练索引的方法，按这个顺序走一遍，基本能覆盖日常用到的所有模式。
第一步，拿一个形状(2, 3, 4)的张量，把x[0]、x[:, 1]、x[..., 2]、x[0, :, 1]、x[:, None]、x[None, ..., 0]逐个打印形状，和预期对照。这一步练的是"整数降维、切片保维、None加维"三条基本规则。
第二步，构造两个能广播的索引张量，比如(2,1)和(3,)，用x[rows, cols]取一次，把结果和手算的配对结果对照。这一步练的是高级索引的广播配对。
第三步，写一个"按样本取top-k"的小函数，要求不用循环，只用索引。提示：topk返回的indices形状是(B, k)，配合arange(B)[:, None]做配对索引。这个练习把高级索引、广播、None补维全串起来了，写完基本就通了。
第四步，把gather、scatter、index_select各写一个最小例子，重点观察输出形状和索引形状的关系。特别是gather，一定要亲手验证"输出形状等于索引形状"这条。
练完这四步，再回头看模型代码里的索引，基本能做到"看到就知道形状怎么变"。索引的难点从来不是语法，而是在脑子里维护形状的变换轨迹——每写一个索引，先想清楚输出形状是什么，再想内容对不对，养成这个习惯，索引相关的bug会少一大半。
