binary_op: dtype mismatch

主线位置：训练主线的「loss」一步。模型输出是 F32，标签从别处来，两者第一次相遇的那个减法。

版本：PyTorch 2.15.0（CPU）、burn 0.22.1（Flex 和 ndarray 设备，debug 和 release 构建）

flowchart TB
  pred["pred F32"] --> sub["pred - target"]
  tgt["target F64"] --> sub
  sub -->|PyTorch| up["提升成 float64"]
  sub -->|"burn debug"| p1["binary_op 断言"]
  sub -->|"burn release"| p2["storage 断言"]
  tgt -->|"cast(pred.dtype())"| fix["F32 再相减"]

实验 crate pita-t 的 p04，PyTorch 侧 p04_promotion.py。平台 M2 Max、macOS 27、rustc 1.99.0。panic 行里的本地路径前缀换成了 ...，note: run with RUST_BACKTRACE=1 一行略去。

PyTorch 的写法

标签用 numpy 读进来，默认是 float64；模型输出是 float32。直接算 loss：

pred = torch.tensor([0.5, 1.5])                  # model output, float32
target = torch.from_numpy(np.array([1.0, 1.0]))  # labels from numpy, float64
print("pred - target ->", (pred - target).dtype)
print("mse_loss      ->", F.mse_loss(pred, target).dtype, F.mse_loss(pred, target).item())
ids = torch.tensor([1, 2], dtype=torch.int32)
print("int32 + int64 ->", (ids + torch.arange(2)).dtype)

pred torch.float32 | target torch.float64
pred - target -> torch.float64
mse_loss      -> torch.float64 0.25
int32 + int64 -> torch.int64
Linear(f64 input): mat1 and mat2 must have the same dtype, but got Double and Float

逐元素运算按类型提升（type promotion）规则选一个更宽的类型，loss 变成 float64，训练照常进行。PyTorch 也不是处处提升：最后一行是把 float64 的输入喂给 float32 的 nn.Linear，矩阵乘法报错。所以 PyTorch 用户通常只在 Linear 那里撞到 dtype 问题，loss 里的混用一直没人发现。

搬到 burn 的现象

标签按原样保留 f64（为了「不丢精度」把 DType::F64 传进 from_data）：

let pred = Tensor::<1>::from_floats([0.5, 1.5], &device);
let target = Tensor::<1>::from_data([1.0f64, 1.0], (&device, DType::F64));
let loss = (pred - target).powi_scalar(2).mean();

### debug
pred F32 | target F64

thread 'main' (8782845) panicked at .../burn-flex-0.22.1/src/ops/binary.rs:41:5:
assertion `left == right` failed: binary_op: dtype mismatch
  left: F32
 right: F64
### release
pred F32 | target F64

thread 'main' (8783081) panicked at .../burn-flex-0.22.1/src/tensor.rs:154:9:
storage: dtype mismatch (expected F64, got F32)

整数也一样，而且不用自己写 dtype 就能碰上：在 Flex 上建的 Int（I32）用 to_device 搬到 ndarray，再和 ndarray 上新建的 Int（I64）相加：

let nd = Device::ndarray();
let ids = Tensor::<1, Int>::from_ints([1, 2], &device).to_device(&nd);
let offs = Tensor::<1, Int>::arange(0..2, &nd);
println!("ids + offs {:?}", iv(ids + offs));

ids I32 on Device<NdArray(Cpu)> | offs I64

thread 'main' (8782877) panicked at .../burn-ndarray-0.22.1/src/ops/int_tensor.rs:152:9:
Data type mismatch (lhs: I32, rhs: I64)

debug 和 release 下这一条相同。设备检查通过了（两个都在 ndarray 上），dtype 不一样。

原因

a - b 进 burn 时先过 TensorCheck::binary_ops_ew，它只查设备和形状，不查 dtype。dtype 由后端自己把关，各后端写法不同：

• Flex 的 binary_op 开头是 debug_assert_eq!，release 下这一行不存在；panic 改由读存储的 FlexTensor::storage 抛出，位置和措辞都变了。那条消息里 expected 后面打印的是张量自己的 dtype，got 是代码想按哪种类型读，和直觉相反，dtype 与设备的「没有类型提升」一节逐字分析过。
• ndarray 用普通的 panic!，两种构建一样。

to_device 只换设备、不换 dtype；dtype 在创建时由当时的设备定下，之后跟着张量走（同一页的「动脑筋答案」）。Flex 默认 I32、ndarray 默认 I64，混用两个设备时就会出现 I32 + I64。

origin/main（2026-10-28）上，#6225「validate device and dtype compatibility across module operations」给 linear、卷积这类模块函数加了 dtype 检查（binary_ops_float_dtype），还没有进任何发布版；逐元素运算的 binary_ops_ew 在 main 上仍然只查设备和形状。

改法

在数据进入模型的地方把 dtype 对齐，不要指望运算时自动处理。cast 接受运行期的 DType，可以直接对齐到另一个操作数：

let target = target.cast(pred.dtype());
let loss = (pred - target).powi_scalar(2).mean();

pred F32 | target F64
target after cast F32
loss [0.25]

和 PyTorch 的 0.25 相同（(0.5-1)² + (1.5-1)² 的平均，f32 和 f64 都精确）。不同之处是 burn 这里在 f32 里算，PyTorch 那边在 float64 里算；如果确实要 f64 的 loss，就把 pred 也 cast(FloatDType::F64)，让两边一起变宽。

整数的情况：混用设备时，to_device 之后显式 cast(IntDType::I64)，或者创建时就写出 dtype。

怎么检查

• 在数据管道的出口断言 dtype：assert_eq!(target.dtype(), DType::F32)。dtype 是运行期的值，不在类型里，编译器帮不上忙。
• 先跑 debug 构建：Flex 的 debug 消息直接写着 binary_op: dtype mismatch 和左右两边的 dtype；release 下的 storage: dtype mismatch 指向读存储的位置，离出错的那一行更远。
• 从 numpy 或文件读来的 f64 数据，不写 dtype 时 from_data 会转成设备默认的 F32（dtype 与设备）；会出问题的往往是你特意写了 DType::F64 的那一处。float64 的 checkpoint 也是同一类问题，见 f64-checkpoint。

小结

• burn 没有类型提升：F32 和 F64、I32 和 I64 相遇就 panic，PyTorch 会静静提升。
• Flex 的报错在 debug 下是 binary_op: dtype mismatch，release 下变成 storage: dtype mismatch，措辞里的 expected / got 是反的。
• to_device 保留原 dtype，Flex 的 I32 搬到 ndarray 后和 I64 相加会 panic。
• 在数据入口 cast(pred.dtype()) 或 cast(FloatDType::F32) 对齐 dtype。