切片与索引
主线位置：张量主线的「计算」一步。切片和掩码是
a[...]这一类读写，发生在张量建好之后、读回主机之前。
版本：PyTorch 2.15.0（CPU）、burn 0.22.1（Flex 设备）
flowchart TB pt["PyTorch a 方括号"] --> rng["范围 1:3"] pt --> idx["整数下标 1"] pt --> fancy["整数张量下标"] pt --> mask["布尔掩码"] rng --> sl["slice s! 保秩"] idx --> sel["select_dim 降秩"] idx --> sl fancy --> select["select"] mask --> msel["mask_select 读"] mask --> mw["mask_where 和 mask_fill 写"]
两边的实验程序对同一个 3×4 的整数张量 arange(12).reshape(3, 4) 做同一组切片，每行打印「名字 形状 值」。PyTorch 侧是 slicing.py，burn 侧是 crate tensorb-slicing（burn =0.22.1，feature std、flex，debug 构建）。平台 M2 Max、macOS 27、rustc 1.99.0。整数结果两边逐个相同，没有误差可言。
范围切片：s! 宏
a = torch.arange(12).reshape(3, 4)
show("cols_1_3", a[:, 1:3])
show("step_2", a[:, ::2])
show("oob", a[:, 2:10])
show("empty_desc", a[:, 3:1])
let a = Tensor::<1, Int>::arange(0..12, &device).reshape([3, 4]);
show("cols_1_3", a.clone().slice(s![.., 1..3]));
show("step_2", a.clone().slice(s![.., ..;2]));
show("oob", a.clone().slice(s![.., 2..10]));
show("empty_desc", a.slice(s![.., 3..1]));
# PyTorch
cols_1_3 [3, 2] [1, 2, 5, 6, 9, 10]
step_2 [3, 2] [0, 2, 4, 6, 8, 10]
oob [3, 2] [2, 3, 6, 7, 10, 11]
empty_desc [3, 0] []
# burn
cols_1_3 [3, 2] [1, 2, 5, 6, 9, 10]
step_2 [3, 2] [0, 2, 4, 6, 8, 10]
oob [3, 2] [2, 3, 6, 7, 10, 11]
empty_desc [3, 0] []
写法一一对应：: 是 ..，1:3 是 1..3，::2 是 ..;2（步长写在分号后面）。不用步长时也可以直接传数组 slice([0..3, 1..3])。Tensor::slice 按值拿走 self，同一个张量要切多次就得先 clone()，原因见运算消费张量。
越界的范围两边都截到形状以内，不报错。burn 截断的位置在 Slice::to_range，机制和它在窗口切分里造成的问题见 burn 卷的切片越界不报错。起点大于终点的范围（3..1）两边都得到长度为 0 的一维。
Tensor::slice 的文档还说降序范围会 panic0.22.1 里 Tensor::slice 的文档「Panics」一节仍写着「If a range is descending (e.g., 2..1) or empty (e.g., 1..1) without negative step」，上面的 3..1 却正常返回了 [3, 0]。改变行为的是 #4083「Fix Slice operation to handle empty ranges」，首次发布于 v0.20.0-pre.4。提交说明写的是让后端处理 start >= end 的空切片，「remove panics for empty or descending slices to match PyTorch and ONNX semantics」。现在 TensorCheck::slice 里留着一段注释说空切片是允许的，理由是和 PyTorch 一致、ONNX 的动态切片范围在运行期可能变空。只有文档那一行没跟着改，main 上也还在。
整数下标：burn 不降秩
show("row_1", a[1])
show("row_1_keep", a[1:2])
show("last_col", a[:, -1])
show("last_col_keep", a[:, -1:])
show("row_1", a.clone().slice(s![1]));
show("row_1_select", a.clone().select_dim::<1>(0, 1));
show("last_col", a.clone().slice(s![.., -1]));
# PyTorch
row_1 [4] [4, 5, 6, 7]
row_1_keep [1, 4] [4, 5, 6, 7]
last_col [3] [3, 7, 11]
last_col_keep [3, 1] [3, 7, 11]
# burn
row_1 [1, 4] [4, 5, 6, 7]
row_1_select [4] [4, 5, 6, 7]
last_col [3, 1] [3, 7, 11]
PyTorch 的 a[1] 去掉一维，burn 的 s![1] 不去。s! 里的整数经 Slice::index 变成范围 1..2，-1 变成「从倒数第一个到末尾」；slice 的返回类型是 Self，秩写在类型里，切完还是 Tensor<2, Int>。按 PyTorch 的习惯把结果标成一维，编译器直接拒绝：
let row: Tensor<1, Int> = a.slice(s![1]);
error[E0308]: mismatched types
--> src/bin/rankerr.rs:5:29
|
5 | let row: Tensor<1, Int> = a.slice(s![1]);
| -------------- ^^^^^^^^^^^^^^ expected `1`, found `2`
| |
| expected due to this
|
= note: expected struct `Tensor<1, _>`
found struct `Tensor<2, _>`
要降秩就用 select_dim，目标秩写在 turbofish 里；它的实现就是 slice_dim(dim, index).squeeze_dim(dim)。负下标在 slice 和 select_dim 里都从末尾数，和 PyTorch 一样。
越界的单个下标不报 IndexError
范围越界两边都截断，单个下标越界就不一样了：
attempt("index_oob", lambda: a[5])
show("index_oob", a.clone().slice(s![5]));
// 另一次运行：
show("select_dim_oob", a.select_dim::<1>(0, 5));
# PyTorch
index_oob IndexError: index 5 is out of bounds for dimension 0 with size 3
# burn
index_oob [0, 4] []
thread 'main' (8536873) panicked at .../burn-tensor-0.22.1/src/tensor/api/base.rs:815:9:
=== Tensor Operation Error ===
Operation: 'Squeeze'
Reason:
1. Can't squeeze dimension 0 because its size is not 1
（panic 行里的本地路径前缀换成了 ...。）
s![5] 是范围 5..6，截到 3..3，得到一个空张量，后面的代码照常往下走。select_dim 会 panic，但报的是 Squeeze：截断后的那一维长度是 0，squeeze_dim 拒绝了它，报错里没有下标 5。下标来自数据（标签、词表 id、batch 里的位置）时，越界在 PyTorch 里当场报错，在 burn 里要自己先断言 index < dims[d]；按单个下标取一行还可以用 narrow(dim, index, 1)，它检查原始下标，见切片越界不报错的修法。
负步长：能用，但含义不同
attempt("reverse", lambda: a[:, ::-1])
show("flip", a.flip(1))
show("reverse", a.clone().slice(s![.., ..;-1]));
show("flip", a.clone().flip([1]));
# PyTorch
reverse ValueError: step must be greater than zero
flip [3, 4] [3, 2, 1, 0, 7, 6, 5, 4, 11, 10, 9, 8]
# burn
reverse [3, 4] [3, 2, 1, 0, 7, 6, 5, 4, 11, 10, 9, 8]
flip [3, 4] [3, 2, 1, 0, 7, 6, 5, 4, 11, 10, 9, 8]
PyTorch 的切片不收负步长，要反转得用 flip。burn 两种都行，slice 的文档写着「Negative steps reverse the selection order」。在 Flex 上两者都只改 Layout（负步长），不复制数据，见 burn 卷 Flex 的「Layout」一节。
a 还是那个 3×4 的 arange(12)，a.slice(s![.., 1..;-1]) 的形状和值是什么？平台、构建同上。提示：NumPy 里 a[:, 1::-1] 是从下标 1 往回走。
整数张量下标与布尔掩码
show("fancy_rows", a[torch.tensor([2, 0])])
show("mask_select", a[a % 2 == 0])
show("where", torch.where(a % 2 == 0, a, -1))
b = a.clone()
b[a % 2 == 0] = 0
show("masked_assign", b)
let idx = Tensor::<1, Int>::from_ints([2, 0], &device);
show("fancy_rows", a.clone().select(0, idx));
let even = a.clone().remainder_scalar(2).equal_elem(0);
show("mask_select", a.clone().mask_select(even.clone()));
let minus_one = Tensor::<2, Int>::full([3, 4], -1, &device);
show("where", minus_one.mask_where(even.clone(), a.clone()));
show("masked_assign", a.clone().mask_fill(even, 0));
# PyTorch
fancy_rows [2, 4] [8, 9, 10, 11, 0, 1, 2, 3]
mask_select [6] [0, 2, 4, 6, 8, 10]
where [3, 4] [0, -1, 2, -1, 4, -1, 6, -1, 8, -1, 10, -1]
masked_assign [3, 4] [0, 1, 0, 3, 0, 5, 0, 7, 0, 9, 0, 11]
# burn
fancy_rows [2, 4] [8, 9, 10, 11, 0, 1, 2, 3]
mask_select [6] [0, 2, 4, 6, 8, 10]
where [3, 4] [0, -1, 2, -1, 4, -1, 6, -1, 8, -1, 10, -1]
masked_assign [3, 4] [0, 1, 0, 3, 0, 5, 0, 7, 0, 9, 0, 11]
| PyTorch | burn | 注意 |
|---|---|---|
a[idx]、index_select | select(dim, idx) | 下标是 Tensor<1, Int>，要写出维度 |
a[mask]、masked_select | mask_select(mask) | 返回 Tensor<1, K>；掩码不广播，形状必须相同 |
torch.where(mask, x, y) | y.mask_where(mask, x) | 接收者是「掩码为假时」的那个张量 |
a[mask] = v、masked_fill | a.mask_fill(mask, v) | 返回新张量 |
mask_where 的参数顺序和 torch.where 相反：torch.where(cond, x, y) 是「条件真取 x」，burn 是「在 y 上，把条件真的位置换成 x」。写反了不会报错，结果正好互补。
mask_select 的结果长度取决于掩码里有几个 true，它的文档说这会让主机和设备同步一次：没有原生实现的后端要把掩码读回主机算下标。在 GPU 设备上，它和 PyTorch 的 a[mask] 一样不适合放在热循环里。
切片不是可写的视图
PyTorch 的基本切片是视图，改视图就是改原张量：
v = a[:, 1:3]
v += 100
show("a_after_view_add", a)
let v = a.clone().slice(s![.., 1..3]);
println!("a.can_mut() while slice alive: {}", a.can_mut());
let v = v + 100;
show("a_after_slice_add", a.clone());
show("slice_assign", a.slice_assign(s![.., 1..3], v));
# PyTorch
a_after_view_add [3, 4] [0, 101, 102, 3, 4, 105, 106, 7, 8, 109, 110, 11]
# burn
a.can_mut() while slice alive: false
a_after_slice_add [3, 4] [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
slice_assign [3, 4] [0, 101, 102, 3, 4, 105, 106, 7, 8, 109, 110, 11]
在 Flex 上，切片其实也和原张量共用缓冲：Flex 的 slice 只是克隆 Arc、换一个 Layout。所以切片活着的时候 a.can_mut() 是 false。共用缓冲却不会出现 PyTorch 那种写穿：v + 100 发现缓冲不是独占的，就另开一块写结果，a 不变。要把结果写回原张量的那一块，用 slice_assign（常数用 slice_fill），它返回新的张量。a 独占缓冲时，Flex 直接在原缓冲上写；这一点和 PyTorch 的原地操作有什么不同，见原地操作。
动脑筋答案
quiz [3, 3] [3, 2, 1, 7, 6, 5, 11, 10, 9]
形状 [3, 3]，每行是下标 3、2、1。burn 先把 1.. 截成范围 1..4，再按负步长倒着取这个范围，等于 a[:, 1:].flip(1)。NumPy 的 a[:, 1::-1] 是从下标 1 开始往回走，得到 [1, 0] 两列：
numpy a[:, 1::-1] (3, 2) [1, 0, 5, 4, 9, 8]
torch flip of a[:, 1:] [[3, 2, 1], [7, 6, 5], [11, 10, 9]]
从 NumPy 搬带负步长的切片时，不能照抄起止下标：burn 的范围永远是「从小到大的那一段」，负号只决定读的方向。
小结
a[:, 1:3]对slice(s![.., 1..3])，步长写成..;2；越界范围两边都截断，降序范围两边都得到空维（文档里「会 panic」那句已过时）。- 整数下标在 burn 里不降秩：
s![1]得到[1, 4]；要[4]用select_dim::<1>(0, 1)。 - 越界的单个下标在 PyTorch 是
IndexError，在 burn 的slice里是空张量，select_dim报的是Squeeze。 - 负步长 burn 能用，含义是「倒着读截好的范围」，和 NumPy 的「从起点往回走」不同。
select、mask_select、mask_where、mask_fill对应整数下标和布尔掩码；mask_where的接收者是条件为假时的值。- 切片在 Flex 上共用缓冲但不写穿；写回用
slice_assign/slice_fill。