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, ...>: ElementConversionis 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 逐位相同。