Rust in Everything
基于 burn 0.22.1

Tensor<1>: ElementConversion is not satisfied

主线位置:训练主线的 forward:手写一层仿射变换 x @ W + b,矩阵加一维偏置。

版本:PyTorch 2.15.0(CPU)、burn 0.22.1(Flex 设备)
flowchart TB
  add["Tensor D2 + Tensor D1"] --> s1["Add Self 要同秩"]
  add --> s2["Add E 要标量"]
  s1 --> no1["不匹配"]
  s2 --> no2["Tensor D1 不是元素"]
  no2 --> err["E0277 ElementConversion"]
  err --> fix["b.unsqueeze()"]

实验 crate pita-t 的 p06、编不过的例子 pita-ce 的 ce06、ce06b,PyTorch 侧 p06_rank_add.py。debug 构建。平台 M2 Max、macOS 27、rustc 1.99.0。

PyTorch 的写法

x = torch.ones(2, 3)
W = torch.ones(3, 4)
b = torch.tensor([0.1, 0.2, 0.3, 0.4])
out = x @ W + b
x @ W + b [2, 4] [3.0999999046325684, 3.200000047683716, 3.299999952316284, 3.4000000953674316]

[2, 4] + [4]:PyTorch 在较短的形状前面补 1,再广播。同样的习惯还有 out * keep(keep 是布尔掩码)和 out / out.mean()(mean() 是 0 维张量)。

搬到 burn 的现象

fn affine(x: Tensor<2>, w: Tensor<2>, b: Tensor<1>) -> Tensor<2> {
  x.matmul(w) + b
}
error[E0277]: the trait bound `Tensor<1>: ElementConversion` is not satisfied
 --> src/bin/ce06.rs:4:15
  |
4 |   x.matmul(w) + b
  |               ^ 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>>`

报错里没有「秩」也没有「形状」,列出来的是 f32、i32 这些元素类型。种类不同时是同一条,乘一个 Bool 掩码:

let keep = Tensor::<2, Bool>::from_bool([[true, false, true], [true, true, false]], &device);
let y = x * keep;
error[E0277]: the trait bound `Tensor<2, burn::tensor::Bool>: ElementConversion` is not satisfied
 --> src/bin/ce06b.rs:7:13
  |
7 |   let y = x * keep;
  |             ^ the trait `ElementConversion` is not implemented for `Tensor<2, burn::tensor::Bool>`
...
  = note: required for `Tensor<2>` to implement `std::ops::Mul<Tensor<2, burn::tensor::Bool>>`

原因

Tensor<D, K> 上的 + 有两个实现:Add<Self> 加一个同秩、同种类的张量,Add<E: ElementConversion> 加一个标量。Tensor<1> 不是 Self(Tensor<2>),编译器只剩标量那一个可选,于是报「Tensor<1> 不是一种元素类型」。*、-、/ 的结构相同。秩在类型里,补维是调用方的事;这一节的来龙去脉见归约与广播的「广播」一节。

burn 自己的 Linear 怎么加偏置:Linear::forward 调函数式的 module::linear,里面写的是 output + bias.unsqueeze()。

改法

把一维的偏置补成 [1, 4],剩下的交给广播:

fn affine(x: Tensor<2>, w: Tensor<2>, b: Tensor<1>) -> Tensor<2> {
  x.matmul(w) + b.unsqueeze()
}
// 或者直接用函数式的 linear,它接受一维偏置
let f = burn::tensor::module::linear(Tensor::<2>::ones([2, 3], &device), w, Some(b));
// Bool 掩码先转成浮点;mean() 的 [1] 先补成 [1, 1]
println!("out * keep.float()          {:?}", fv(out.clone() * keep.float()));
let scale = out.clone().mean(); // Tensor<1>, shape [1]
println!("out / out.mean().unsqueeze() {:?}", fv((out / scale.unsqueeze()).slice([0..1])));
x.matmul(w) + b.unsqueeze() [2, 4] [3.1, 3.2, 3.3, 3.4]
module::linear(x, w, Some(b))  [3.1, 3.2, 3.3, 3.4]
out * keep.float()          [3.1, 0.0, 3.3, 3.4, 3.1, 3.2, 0.0, 3.4]
out / out.mean().unsqueeze() [0.9538461, 0.9846154, 1.0153846, 1.0461539]

PyTorch 侧对应的三行:

x @ W + b [2, 4] [3.0999999046325684, 3.200000047683716, 3.299999952316284, 3.4000000953674316]
out * keep [3.0999999046325684, 0.0, 3.299999952316284, 3.4000000953674316, 3.0999999046325684, 3.200000047683716, 0.0, 3.4000000953674316]
out / out.mean() [0.9538460969924927, 0.9846153855323792, 1.015384554862976, 1.0461539030075073] | out.mean() shape []

数字看起来不同,其实是同一组 f32:burn 打印 f32 的最短表示,Python 的 tolist() 打印换成 f64 之后的全部位数。把两边的字符串各自转成 np.float32 逐个比较,仿射、掩码、除法三组共 16 个数全部相等,差为 0。

unsqueeze::<2>() 补在最前面,[4] 变 [1, 4],对应 PyTorch 自动补的那一维。要补在后面([B] 变 [B, 1])用 unsqueeze_dim::<2>(1),两者不能互换,见 keepdim 那个算错的准确率。into_scalar() 也能把 mean() 变成一个 Rust 的数再除,但它要把数据读回主机,在 GPU 上会多一次同步。

怎么检查

  • 看到「the trait bound Tensor<N, ...>: ElementConversion is not satisfied」,不要去找元素类型,先看运算符两边的秩和种类(Float、Int、Bool)是否一样。
  • 秩不同:unsqueeze / unsqueeze_dim 补维;种类不同:.float()、.int()、.bool() 换种类。
  • 自己写的层尽量复用 burn::tensor::module::linear 这类函数式接口,它们已经处理了偏置的秩。

小结

  • Tensor<2> + Tensor<1> 编译不过,E0277 说的是 Tensor<1>: ElementConversion,因为 + 只剩加标量的那个重载可选。
  • 种类不同(乘 Bool 掩码)报同一条。
  • 偏置用 b.unsqueeze(),掩码用 keep.float(),全局归约的 [1] 用 unsqueeze() 补成同秩;结果和 PyTorch 逐位相同。

💬 关联讨论

发起讨论
加载中...