argmax 是 [B, 1]，准确率算成 0.4375

主线位置：训练主线的 forward 之后：拿 logits 算一个指标（准确率），用到按维归约。

版本：PyTorch 2.15.0（CPU）、burn 0.22.1（Flex 设备）

flowchart TB
  lg["logits 4x3"] --> am["argmax 1"]
  am -->|PyTorch| p1["pred 形状 4"]
  am -->|burn| b1["pred 形状 4x1"]
  b1 -->|"y.unsqueeze()"| bad["广播成 4x4"]
  bad --> wrong["acc 0.4375"]
  b1 -->|"squeeze_dim 1"| ok["acc 0.75"]

实验 crate pita-t 的 p05、编不过的例子 pita-ce 的 ce05，PyTorch 侧 p05_keepdim.py。debug 构建。平台 M2 Max、macOS 27、rustc 1.99.0。

PyTorch 的写法

logits = torch.tensor([[2.0, 0.1, 0.3], [0.2, 1.5, 0.1], [0.1, 0.2, 3.0], [1.0, 2.0, 0.5]])
y = torch.tensor([0, 1, 1, 1])
pred = logits.argmax(dim=1)
acc = (pred == y).float().mean()

pred [4] [0, 1, 2, 1]
acc  [] 0.75

argmax(dim=1) 默认 keepdim=False，pred 和 y 都是 [4]，四个里对了三个。

搬到 burn 的现象

逐字翻译，编译器先拦住：

let acc = logits.argmax(1).equal(y).float().mean();

error[E0308]: mismatched types
    --> src/bin/ce05.rs:7:36
     |
   7 |   let acc = logits.argmax(1).equal(y).float().mean();
     |                              ----- ^ expected `2`, found `1`
     |                              |
     |                              arguments to this method are incorrect
     |
     = note: expected struct `Tensor<2, _>`
                found struct `Tensor<1, _>`
note: method defined here
    --> .../burn-tensor-0.22.1/src/tensor/api/base.rs:2464:12
     |
2464 |     pub fn equal(self, other: Self) -> Tensor<D, Bool> {
     |            ^^^^^

（note 里的本地路径前缀换成了 ...。）argmax(1) 的结果是 Tensor<2, Int>，y 是 Tensor<1, Int>。编译器要一个二维的 y，最顺手的改法是 y.unsqueeze()，于是编过了：

let pred = logits.argmax(1);
let hit = pred.clone().equal(y.clone().unsqueeze());
let acc = hit.float().mean();

pred [4, 1] [0, 1, 2, 1]
pred.equal(y.unsqueeze()) [4, 4]
acc  [1] [0.4375]

准确率变成 0.4375，没有任何报错。unsqueeze() 把新维度加在前面，y 成了 [1, 4]；[4, 1] 和 [1, 4] 比较时两边都广播，得到一个 [4, 4] 的矩阵，算的是「每个预测和每个标签」是否相等。预测为 1 的两行各和三个标签 1 相等，预测为 0 的一行和一个标签相等，一共 7 个，7 / 16 = 0.4375。

原因

burn 的按维归约没有 keepdim 参数，永远保留被归约的维度：argmax 返回 Tensor<D, Int>，sum_dim 返回 Self，秩不变；为什么稳定版 Rust 写不出「秩减一」的返回类型，见归约与广播的第一节。全局的 sum()、mean() 返回 [1] 而不是 0 维，同一页的第二节。

秩对不上时编译器会报错，形状对不上却未必。equal 和加减乘除一样经过 binary_ops_ew_shape，按广播规则检查：每一维要么相等、要么一边是 1。[4, 1] 对 [1, 4] 完全合法。这和 PyTorch 一致，在上面的 PyTorch 脚本里加一行 pred[:, None] == y[None, :]，形状也是 [4, 4]；区别只在 burn 先多给了你一个大小为 1 的维度。

改法

把保留下来的那一维去掉，再和 [B] 的标签比：

let hit = pred.squeeze_dim::<1>(1).equal(y);
let acc: f32 = hit.float().mean().into_scalar();

另一种写法是把标签补成 [B, 1]，新维度加在第 1 维而不是第 0 维：

let alt: f32 = pred.clone().equal(y.clone().unsqueeze_dim::<2>(1)).float().mean().into_scalar();

y.unsqueeze_dim(1): acc 0.75
squeeze_dim(1): acc 0.75

两种都和 PyTorch 的 0.75 相同（整数比较，结果是 3 / 4，没有误差）。squeeze_dim 要求那一维的大小是 1，否则 panic，比 squeeze() 安全：squeeze() 去掉所有大小为 1 的维，batch 为 1 时会把 batch 维也去掉。

mean() 是 [1]，要一个 Rust 的数就 into_scalar()，对应 .item()。

怎么检查

• 每次 sum_dim、mean_dim、max_dim、argmax、argmin 之后，问自己：PyTorch 版本里这里写没写 keepdim=True？没写的，在 burn 里紧跟一个 squeeze_dim。
• 编译器要求「补一维」时，先想清楚补在哪一边。unsqueeze() 补在最前面，unsqueeze_dim(d) 补在第 d 维；两者在方阵、batch 等于类别数之类的情形下都能编译，结果不同。
• 对指标和 loss 的中间结果打印一次 dims()。逐元素比较的结果形状应当和输入一样，变大了就是广播了。

小结

• burn 的 argmax、sum_dim 等按维归约保留这一维，相当于 PyTorch 的 keepdim=True；全局归约返回 [1]。
• [B, 1] 和 [B] 比较编译不过；按 unsqueeze() 补成 [1, B] 能编译，广播成 [B, B]，准确率静默算错。
• 先 squeeze_dim::<1>(1) 再比较，或者 unsqueeze_dim::<2>(1) 补在后面。