Rust in Everything
基于 burn 0.22.1

归约与广播

主线位置：张量主线的中间一步「计算」：按维归约、全局归约，以及两个形状不同的张量怎么做逐元素运算。
版本：PyTorch 2.15.0（CPU）、burn 0.22.1（Flex 设备）
flowchart TB
  x["a 形状 2x3"] -->|"sum_dim 1"| k["形状 2x1"]
  k -->|"squeeze_dim 1"| s["Tensor 1 形状 2"]
  x -->|"sum()"| one["Tensor 1 形状 1"]
  one -->|into_scalar| f["f32 标量"]
  k -->|广播| sub["a 减 k 形状 2x3"]
  row["Tensor 1 形状 3"] -->|unsqueeze| r2["形状 1x3"]
  r2 -->|广播| sub

本页的对照：PyTorch 侧 tensora_reduce.py，burn 侧实验 crate tensora-t 的 reduce（feature 同创建与形状；debug 构建）。PyTorch 侧 torch.set_num_threads(1)。平台 M2 Max、macOS 27、rustc 1.99.0。panic 信息里 cargo 缓存的路径前缀用 ... 略去。两边的输入都是 arange(6) 排成的 [2, 3] 浮点矩阵 a。

按维归约：burn 保留这一维

a = torch.arange(6, dtype=torch.float32).reshape(2, 3)
a.sum(dim=1); a.sum(dim=1, keepdim=True); a.sum(dim=0)
a.mean(dim=1)

let a = Tensor::<1, Int>::arange(0..6, &device).float().reshape([2, 3]);
a.clone().sum_dim(1); a.clone().sum_dim(0); a.clone().sum_dim(-1);
a.clone().sum_dim(1).squeeze_dim::<1>(1);
a.clone().sum_dim(1).squeeze::<1>();
a.clone().mean_dim(1);

sum(dim=1)             [2] float32 [3.0, 12.0]
sum(dim=1, keepdim)    [2, 1] float32 [3.0, 12.0]
sum(dim=0)             [3] float32 [3.0, 5.0, 7.0]
mean(dim=1)            [2] float32 [1.0, 4.0]

sum_dim(1)             [2, 1] [3.0, 12.0]
sum_dim(0)             [1, 3] [3.0, 5.0, 7.0]
sum_dim(-1)            [2, 1] [3.0, 12.0]
sum_dim(1).squeeze_dim [2] [3.0, 12.0]
sum_dim(1).squeeze()   [2] [3.0, 12.0]
mean_dim(1)            [2, 1] [1.0, 4.0]

数值一样，形状不一样。PyTorch 的默认是 keepdim=False，被归约的维度消失；burn 的 sum_dim、mean_dim、max_dim、argmax 没有这个参数，永远相当于 keepdim=True。负下标两边都支持。

原因在签名上：pub fn sum_dim<I: AsIndex>(self, dim: I) -> Self。返回 Self 就是同一个 Tensor<D, K>，秩不变。要让秩减一，返回类型得写成 Tensor<{ D - 1 }>，稳定版 Rust 不允许：

fn sum_drop<const D: usize>(t: Tensor<D>, dim: usize) -> Tensor<{ D - 1 }> {
  t.sum_dim(dim).squeeze_dim(dim)
}

error: generic parameters may not be used in const operations
 --> src/bin/ce_sumrank.rs:3:67
  |
3 | fn sum_drop<const D: usize>(t: Tensor<D>, dim: usize) -> Tensor<{ D - 1 }> {
  |                                                                   ^ cannot perform const operation using `D`
  |
  = help: const parameters may only be used as standalone arguments here, i.e. `D`

所以去掉这一维要单独一步，并且由调用方写出新的秩：squeeze_dim::<1>(1) 去掉指定的那一维（它的大小必须是 1，否则 panic Can't squeeze dimension 1 because its size is not 1），squeeze::<1>() 去掉所有大小为 1 的维，和 PyTorch 不带参数的 squeeze() 一样，batch 为 1 时会连 batch 维一起去掉。

📝NOTE
冷知识：squeeze 以前是现在的 squeeze_dim

0.19 之前的 burn 代码按下标去掉一维写的是 squeeze::<1>(1)（burn 自己的 burn-nn、burn-train 里也这样写过），在 0.22.1 上编译不过：

error[E0061]: this method takes 0 arguments but 1 argument was supplied
   --> src/bin/ce_squeeze.rs:6:35
    |
  6 |   let s: Tensor<1> = a.sum_dim(1).squeeze::<1>(1);
    |                                   ^^^^^^^^^^^^ - unexpected argument of type `{integer}`
...

#3790「[Breaking] Make squeeze/squeeze_dim consistent with other APIs」（首次发布于 v0.19.0）之前，squeeze::<D2>(dim) 就是按下标去掉一维。这个 PR 把它改名为 squeeze_dim，再新加一个不带参数的 squeeze::<D2>()，去掉所有大小为 1 的维，并把 burn-nn、burn-train、burn-import 里的调用一起改了。旧代码里的 squeeze::<1>(1) 换成 squeeze_dim::<1>(1) 即可。

全局归约返回 [1]，不是标量

a.sum(); a.mean(); a.max(); a.sum().item()

a.clone().sum(); a.clone().mean(); a.clone().max();
let total: f32 = a.clone().sum().into_scalar();

sum()                  [] float32 [15.0]
mean()                 [] float32 [2.5]
max()                  [] float32 [5.0]
sum().item()           15.0

sum()                  [1] [15.0]
mean()                 [1] [2.5]
max()                  [1] [5.0]
sum().into_scalar()    15

sum 和 mean 的返回类型都是 Tensor<1, K>。burn 没有 0 维张量（建 Tensor<0> 会 panic，见创建与形状），一个数只能装进形状 [1]。要拿到 Rust 的数，into_scalar()，对应 .item()。

把这样一个结果当标量参与后面的运算时要注意：PyTorch 的 0 维张量和任何形状都能广播；burn 的 Tensor<1> 要和 Tensor<2> 运算，得先 unsqueeze 成 [1, 1]（见下面广播一节）。

带下标的 max 和 argmax

r = a.max(dim=1)
r.values; r.indices; a.argmax(dim=1)

let (v, i) = a.clone().max_dim_with_indices(1);
a.clone().argmax(1);

max(dim=1) -> max
max(dim=1).values      [2] float32 [2.0, 5.0]
max(dim=1).indices     [2] int64 [2, 2]
argmax(dim=1)          [2] int64 [2, 2]

max_dim_with_indices.0 [2, 1] [2.0, 5.0]
max_dim_with_indices.1 [2, 1] I32 [2, 2]
argmax(1)              [2, 1] I32 [2, 2]

torch.max(dim=...) 返回一个名为 max 的具名元组；burn 分成两个方法：max_dim 只给值，max_dim_with_indices 返回 (值, 下标) 的普通元组。下标张量的种类是 Int，dtype 是设备默认整数（Flex 上 I32，PyTorch 是 int64），形状同样保留这一维。argmax 的返回类型 Tensor<D, Int> 写明了这两点。拿 argmax 和 [batch] 形状的标签比较前，先 squeeze_dim::<1>(1)。

var 的默认校正

a.var(dim=1); a.var(dim=1, correction=0)

a.clone().var(1); a.clone().var_bias(1);

var(dim=1)             [2] float32 [1.0, 1.0]
var(dim=1, correction=0) [2] float32 [0.6666666865348816, 0.6666666865348816]

var(1)                 [2, 1] [1.0, 1.0]
var_bias(1)            [2, 1] [0.6666667, 0.6666667]

两边的 var 默认都带贝塞尔校正（除以 n−1）。不要校正时，PyTorch 传 correction=0，burn 换一个方法 var_bias。两边的 0.6666667 是同一个 f32，PyTorch 的 tolist() 把它按 f64 打印，所以多出一串数字。

整数的归约

i = torch.tensor([1, 2])
i.mean(); i.sum()
torch.tensor([1, 2], dtype=torch.int32).sum()

let i = Tensor::<1, Int>::from_ints([1, 2], &device);
i.clone().mean(); i.sum();

int mean(): NotImplementedError: mean(): could not infer output dtype. Input dtype must be either a floating point or complex dtype. Got: Long
int sum()              [] int64 [3]
int32 sum()            [] int64 [3]

Int mean()             [1] I32 [1]
Int sum()              [1] I32 [3]

两处不同：

  • PyTorch 拒绝对整数求 mean；burn 的 mean 定义在 Numeric 上，Int 也有，结果是整数，1.5 变成了 1。要平均值，先 .float()。
  • PyTorch 对 int32 求和时把结果提升成 int64；burn 的 I32 求和还是 I32，可能溢出，见整数运算。

广播：规则相同，秩要先对齐

两边的广播规则一样：从最后一维往前对齐，每一维要么相等、要么有一边是 1。

row = torch.tensor([10.0, 20.0, 30.0])
col = torch.tensor([100.0, 200.0])
a + row
a + col[:, None]
col[:, None] + row
row.expand(2, 3)
a + torch.ones(2, 2)

let row = Tensor::<1>::from_floats([10.0, 20.0, 30.0], &device);
let col = Tensor::<1>::from_floats([100.0, 200.0], &device);
a.clone() + row.clone().unsqueeze::<2>();
a.clone() + col.clone().unsqueeze_dim::<2>(1);
col.clone().unsqueeze_dim::<2>(1) + row.clone().unsqueeze::<2>();
row.clone().unsqueeze::<2>().expand([2, 3]);
a + Tensor::<2>::ones([2, 2], &device);

a + row                [2, 3] float32 [10.0, 21.0, 32.0, 13.0, 24.0, 35.0]
a + col[:, None]       [2, 3] float32 [100.0, 101.0, 102.0, 203.0, 204.0, 205.0]
[2,1] + [3]            [2, 3] float32 [110.0, 120.0, 130.0, 210.0, 220.0, 230.0]
row.expand(2, 3)       [2, 3] float32 [10.0, 20.0, 30.0, 10.0, 20.0, 30.0]
a + ones(2,2): RuntimeError: The size of tensor a (3) must match the size of tensor b (2) at non-singleton dimension 1

a + row.unsqueeze()    [2, 3] [10.0, 21.0, 32.0, 13.0, 24.0, 35.0]
a + col.unsqueeze_dim(1) [2, 3] [100.0, 101.0, 102.0, 203.0, 204.0, 205.0]
[2,1] + [1,3]          [2, 3] [110.0, 120.0, 130.0, 210.0, 220.0, 230.0]
row.expand([2, 3])     [2, 3] [10.0, 20.0, 30.0, 10.0, 20.0, 30.0]

thread 'main' (8547909) panicked at .../burn-tensor-0.22.1/src/tensor/api/numeric.rs:40:9:
=== Tensor Operation Error ===
  Operation: 'Add'
  Reason:
    1. The provided tensors have incompatible shapes. Incompatible size at dimension '1' => '3 != 2', which can't be broadcasted. Lhs tensor shape Shape { dims: [2, 3] }, Rhs tensor shape Shape { dims: [2, 2] }.

形状对不上时两边都报错，burn 的检查是 TensorCheck 里的 binary_ops_ew_shape。区别在于秩：PyTorch 会在较短的形状前面补 1，[2, 3] + [3] 直接算；burn 的 add 要求两边都是 Tensor<D, K>，秩必须相同，补维要自己写：

忘了补维时，报错不提秩，也不提形状：

let a = Tensor::<2>::zeros([2, 3], &device);
let row = Tensor::<1>::from_floats([10.0, 20.0, 30.0], &device);
let b = a + row;

error[E0277]: the trait bound `Tensor<1>: ElementConversion` is not satisfied
 --> src/bin/ce_bcast.rs:7:13
  |
7 |   let b = a + row;
  |             ^ the trait `ElementConversion` is not implemented for `Tensor<1>`
  |
  = help: the following other types implement trait `ElementConversion`:
            bf16
            bool
            burn::tensor::f16
            f32
            f64
            flex32
            i16
            i32
          and 6 others
  = note: required for `Tensor<2>` to implement `std::ops::Add<Tensor<1>>`

Tensor<D, K> 上有两个 +：Add<Self> 加同类型的张量，Add<E: ElementConversion> 加一个标量。Tensor<1> 不是 Self，编译器就去试标量那一个，最后报「Tensor<1> 不是一种元素类型」。看到这条报错，先看两边的秩。Tensor<1, Int> + Tensor<1>（种类不同）报的也是这一条。

❗IMPORTANT
动脑筋：「每一行减去这一行的均值」，两边各写最顺手的一行，谁会报错？
a - a.mean(dim=1)

a.clone() - a.clone().mean_dim(1)

a 就是上面的 [2, 3] 矩阵。再想一步：如果 a 是 [3, 3] 的方阵呢？

对得上的数字

上面的小矩阵两边逐位相同。换一个大一点的输入：64×1000 的 f32 矩阵，元素是 sin(k)。输入只在 PyTorch 侧生成一次，写进 JSON，burn 侧读同一份，这样比较的只是归约本身，不受两边 sin 实现差一个 ulp 的影响：

k = np.arange(rows * cols, dtype=np.float32)
x = torch.from_numpy(np.sin(k)).reshape(rows, cols)
out = {
  "sum_dim1": x.sum(dim=1).tolist(), "mean_dim0": x.mean(dim=0).tolist(),
  "sum": x.sum().item(), "mean": x.mean().item(), "var_dim1": x.var(dim=1).tolist(),
}

let v: Vec<f32> = serde_json::from_str(&text).unwrap();
let x = Tensor::<2>::from_data(TensorData::new(v, [rows, cols]), &device);
let sum_dim1 = to_vec(x.clone().sum_dim(1));
let mean_dim0 = to_vec(x.clone().mean_dim(0));
let total: f32 = x.clone().sum().into_scalar();
let mean: f32 = x.clone().mean().into_scalar();
let var1 = to_vec(x.clone().var(1));

compare.py，容差 1e-5：

sum_dim1                 n= 64  max|diff|=4.41e-06
mean_dim0                n=1000  max|diff|=1.68e-08
sum                      n=  1  max|diff|=8.58e-06
mean                     n=  1  max|diff|=1.34e-10
var_dim1                 n= 64  max|diff|=7.15e-07
worst=8.58e-06 tol=1e-05 OK

不是 0，是因为 f32 加法不满足结合律，两边把 1000 个数加起来的顺序不同。拿同一份 f32 输入在 f64 里求和当参考，两边离它都有误差，各有远近：

字段PyTorch 对 f64 参考burn 对 f64 参考
sum_dim13.61e-061.90e-06
sum5.41e-063.18e-06
var_dim12.97e-087.08e-07

误差的量级看的是「参与求和的数的绝对值之和」，不是结果本身：每行 1000 个 |sin(k)| 之和约 637，f32 的 ε（1.19e-7）乘上它约 7.6e-5，两边每行和的差 4.41e-6 在这个范围以内。全部 64000 个数的和本身只有 0.374，绝对误差 8.58e-6 看起来就很大。所以这里用 1e-5 而不是整卷默认的 1e-6：1e-6 会让 sum_dim1 和 sum 不通过，而它们的差完全来自求和顺序。

反向验证这个容差：burn 侧把输入的第一个元素加 1e-3 再跑，compare.py 报 MISMATCH，退出码 1：

sum_dim1                 n= 64  max|diff|=1.00e-03
mean_dim0                n=1000  max|diff|=1.56e-05
sum                      n=  1  max|diff|=1.01e-03
mean                     n=  1  max|diff|=1.58e-08
var_dim1                 n= 64  max|diff|=7.15e-07
worst=1.01e-03 tol=1e-05 MISMATCH

动脑筋答案

a - a.mean(dim=1): RuntimeError: The size of tensor a (3) must match the size of tensor b (2) at non-singleton dimension 1

a - a.mean_dim(1)      [2, 3] [-1.0, 0.0, 1.0, -1.0, 0.0, 1.0]

PyTorch 报错，burn 算对了。a.mean(dim=1) 的形状是 [2]，广播时和 a 的最后一维 3 对齐，对不上；PyTorch 里要写 keepdim=True。burn 的 mean_dim(1) 本来就是 [2, 1]，和 [2, 3] 直接广播。

方阵更危险。[3, 3] 时 PyTorch 那一行不报错，[3] 被当成行向量，每一行减去的是各行均值组成的向量：

a - a.mean(dim=1)          [[-1.0, -3.0, -5.0], [2.0, 0.0, -2.0], [5.0, 3.0, 1.0]]
a - a.mean(dim=1, keepdim) [[-1.0, 0.0, 1.0], [-1.0, 0.0, 1.0], [-1.0, 0.0, 1.0]]

s - s.mean_dim(1)      [3, 3] [-1.0, 0.0, 1.0, -1.0, 0.0, 1.0, -1.0, 0.0, 1.0]

burn 保留维度的默认值在这里正好挡住了这个错误。反过来，从 PyTorch 搬代码时，凡是原代码没写 keepdim=True、后面又依赖结果是 [2] 的地方，在 burn 里都要补一个 squeeze_dim。

小结

  • sum_dim、mean_dim、max_dim、argmax 永远保留被归约的维度（相当于 keepdim=True），因为返回类型是同秩的 Self；去掉它用 squeeze_dim::<D2>(dim)。
  • 0.19 起 squeeze::<D2>() 不带参数、去掉所有大小为 1 的维，旧的带下标版本叫 squeeze_dim。
  • 全局 sum、mean、max 返回 Tensor<1>，形状 [1]；取数用 into_scalar。
  • argmax 和 max_dim_with_indices 的下标是设备默认整数（Flex 上 I32）；var 两边默认都除以 n−1，不校正用 var_bias。
  • Int 可以 mean，结果是整数；I32 求和不提升到 I64。
  • 广播规则相同，但秩要先用 unsqueeze / unsqueeze_dim 对齐；秩不同时报错说的是 ElementConversion。
  • 64×1000 的 f32 归约两边最大差 8.58e-6（容差 1e-5），来自求和顺序。

💬 关联讨论

发起讨论
加载中...