Rust in Everything
基于 burn 0.22.1

切片越界不报错

主线位置:算子主线的入口,Tensor 的泛型方法里、离开泛型之前。Tensor::slice 先把切片参数换算成 Vec<Slice>,换算时已经截断,之后的 TensorCheck 看到的永远是合法范围(机制见检查与报错的「切片:越界不报错」)。

flowchart TB
  user["s! start..start+4"] --> conv["into_slices 截到形状内"]
  conv --> check["TensorCheck slice 通过"]
  check --> short["窗口变短或变空"]
  short --> later["stack 报形状不一致"]
  user -.->|"改用 narrow"| narrowck["TensorCheck narrow"]
  narrowck --> panicnow["当场 panic"]

本页实验在 crate exp-burn/pb-pit(bin slice、slice_fix),Device::flex(),debug 构建,M2 Max、macOS 27、rustc 1.99.0。

现象

把一段长 10 的序列按 4 个一组切成窗口,再叠成一个批:

let seq = Tensor::<1, Int>::arange(0..10, &device).float();
let mut windows = Vec::new();
for start in (0..10).step_by(4) {
  let w = seq.clone().slice(s![start..start + 4]);
  println!("window {start}..{}: dims {:?} {:?}", start + 4, w.dims(), w.to_data().to_vec::<f32>().unwrap());
  windows.push(w);
}
let past_end = seq.clone().slice(s![12..16]);
println!("12..16 on [10]: dims {:?}", past_end.dims());
println!("stacking {} windows ...", windows.len());
let batch = Tensor::stack::<2>(windows, 0);
window 0..4: dims [4] [0.0, 1.0, 2.0, 3.0]
window 4..8: dims [4] [4.0, 5.0, 6.0, 7.0]
window 8..12: dims [2] [8.0, 9.0]
12..16 on [10]: dims [0]
stacking 3 windows ...

thread 'main' (3479610) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/burn-tensor-0.22.1/src/tensor/api/base.rs:2630:9:
=== Tensor Operation Error ===
  Operation: 'Stack'
  Reason:
    1. Can't stack tensors with different shapes Provided dimension (0), tensors shapes: [Shape { dims: [4] }, Shape { dims: [4] }, Shape { dims: [2] }] 

note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
exit=101

8..12 在长 10 的维度上得到两个元素,12..16 整个落在外面,得到一个形状 [0] 的空张量,都没有报错。真正 panic 的是 stack,报的是「形状不一致」,不提切片。如果后面不是 stack 而是逐个窗口送进模型,短窗口和空窗口会一路算下去:均值、损失照常出数,只是这几个样本的含义已经变了。

Rust 自己的切片不是这样:&v[8..12] 在长 10 的 Vec 上直接 panic。从标准库习惯过来的人最容易在这里掉进去。

出处

Tensor::slice 的前两行:

        let shape = self.shape();
        let slices = slices.into_slices(&shape);

        // Validate slices
        check!(TensorCheck::slice::<D>(&shape, &slices));

into_slices 对每一维调 slice.to_range(dim_size),后者经 convert_signed_index 换算下标:

fn convert_signed_index(index: isize, size: usize) -> usize {
    if index < 0 {
        (size as isize + index).max(0) as usize
    } else {
        (index as usize).min(size)
    }
}

正下标 min(size),负下标 max(0):两头都截。TensorCheck::slice 里那条「range exceeds the current tensor size」检查因此永远看不到越界的范围。这是有意的设计,Slice 的文档写着「Indices are clamped to valid ranges」;slice 之后如果有一维是 0,直接返回一个空张量,不进后端。

切片里会 panic 的只有切片个数多于秩、步长为 0、slice_assign 的值形状不对这几种,原文在检查与报错里。

修法

要「越界就是错」,用 narrow(dim, start, length)。它在调用 slice 之前先跑 TensorCheck::narrow,检查的是用户给的原始数字:length 为 0、start 超出、start + length 超出各有一条。窗口的个数先按整除算好,尾巴怎么处理写在代码里:

let (len, win) = (seq.dims()[0], 4);
// only full windows; the tail is handled on purpose
let windows: Vec<_> = (0..len / win).map(|i| seq.clone().narrow(0, i * win, win)).collect();
let batch = Tensor::stack::<2>(windows, 0);
println!("batch {:?}, tail of {} dropped", batch.dims(), len % win);
// narrow checks the bounds that slice clamps away
let _ = seq.narrow(0, 8, 4);
batch [2, 4], tail of 2 dropped

thread 'main' (3479645) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/burn-tensor-0.22.1/src/tensor/api/base.rs:2701:9:
=== Tensor Operation Error ===
  Operation: 'Narrow'
  Reason:
    1. Can't narrow at dimension 0, start + length exceeds the size of the tensor along this dimension (Size=10) 

note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
exit=101

整除得到两个完整窗口;最后一行故意用 narrow(0, 8, 4) 越界,当场 panic,位置就是这次调用。narrow 一次只管一维;多维的切片仍然用 slice,越界要自己在调用前 assert!。要按固定大小切完整个维度也可以用 split(4, 0),它的文档写明除不尽时最后一块更小,同样需要自己判断。

怎么确认没中招

  • 搜代码里所有 slice( 和 s![:范围里有变量(start + len、i * step)的地方,要么改成 narrow,要么在前面断言 end <= dims[d]。
  • 切完马上 assert_shape!(宏在 检查与报错 里),把「期望的形状」写进代码,短窗口会在切片后的下一行报出来,而不是在几层之外的 stack 或损失里。
  • 数据管道里的最后一批常常更短:如果下游要求固定形状,打印一次每批的 dims()。

小结

  • slice / s! 把越界的范围截到形状以内(正下标 min(size)、负下标 max(0)),得到更短或为空的张量,不报错。
  • 报错会出现在后面某个要求形状一致的算子上(本例是 stack),信息里没有切片。
  • 要越界即报错,用 narrow:它先检查原始的 start、length,再调 slice。

💬 关联讨论

发起讨论
加载中...