shape mismatch in add / matmul / where_cond

主线位置：推理主线的「建模块」一步。手写注意力和线性层时，加因果掩码、加偏置、乘一个向量。

版本：PyTorch 2.15.0（CPU）、candle 0.11.0（CPU，accelerate；Metal）

flowchart TB
  py["PyTorch 写法"] --> add["scores + mask"]
  py --> mv["W @ v"]
  py --> wh["where 配标量"]
  add --> e1["add 报错"]
  mv --> e2["matmul 报错"]
  wh --> e3["where_cond 报错"]
  e1 --> f1["broadcast_add"]
  e2 --> f2["unsqueeze 再 squeeze"]
  e3 --> f3["broadcast_as"]

本页的对照：PyTorch 侧是实验代码里的 p02.py，candle 侧是实验 crate pita-t 的 p02_bcast（features metal accelerate，release）。平台 M2 Max、macOS 27、rustc 1.99.0。CPU 和 Metal 的输出逐字相同，下面只列 CPU。

PyTorch 的写法

注意力里最常见的四行：

scores = torch.arange(18, dtype=torch.float32).reshape(1, 2, 3, 3)  # [b, h, s, s]
mask = torch.triu(torch.full((3, 3), float("-inf")), diagonal=1)    # [s, s]
scores + mask
x @ w.T + bias                                                      # [2, 3] @ [3, 4] + [4]
w @ v                                                               # [4, 3] @ [3]
torch.where(causal, scores, float("-inf"))                          # bool [s, s], scalar

scores + mask [1, 2, 3, 3] [[0.0, -inf, -inf], [3.0, 4.0, -inf], [6.0, 7.0, 8.0]]
x @ w.T + bias [2, 4]
w @ v [4] [3.0, 12.0, 21.0, 30.0]
torch.where(causal, scores, -inf) [1, 2, 3, 3]

四行都靠隐式规则：+ 按广播把 [s, s] 扩成 [b, h, s, s]；@ 把一维的 v 当成列向量、乘完再去掉那一维；torch.where 把条件和标量一起广播到 scores 的形状。

搬到 candle 的现象

&scores + &mask;
x.matmul(&w.t()?)? + &bias;
w.matmul(&v);
causal.where_cond(&scores, &Tensor::new(f32::NEG_INFINITY, &dev)?);

scores + mask                      Err: shape mismatch in add, lhs: [1, 2, 3, 3], rhs: [3, 3]
x.matmul(w.t()) + bias             Err: shape mismatch in add, lhs: [2, 4], rhs: [4]
w.matmul(v)                        Err: shape mismatch in matmul, lhs: [4, 3], rhs: [3]
causal.where_cond(scores, -inf)    Err: shape mismatch in where_cond, lhs: [3, 3], rhs: [1, 2, 3, 3]

四个都是 Err，不 panic，也不会静默算出别的结果。

原因

• + - * / 和 add、mul 这些方法调 same_shape_binary_op，两边形状必须完全相同；广播是另一组方法 broadcast_add 等。报错里的 op 名就是线索：add 说明漏了 broadcast_。机制见形状与广播的「+ 不广播」和源码卷的 shape mismatch in add。

• Tensor::matmul 要求两边秩相同且至少为 2，没有 PyTorch 对一维操作数的补维规则。

• where_cond 开头连着两次 same_shape_binary_op，先比条件和真分支，再比条件和假分支：

  Copy

      pub fn where_cond(&self, on_true: &Self, on_false: &Self) -> Result<Self> {
          let _shap = self.same_shape_binary_op(on_true, "where_cond")?;
          let shape = self.same_shape_binary_op(on_false, "where_cond")?;
  // ...
  

  所以报错里 lhs 是条件 [3, 3]，rhs 是 scores；把这一处修好之后，0 维的 -inf 还会在第二次检查时再报一次。

改法

scores.broadcast_add(&mask);
x.matmul(&w.t()?)?.broadcast_add(&bias);
w.matmul(&v.unsqueeze(1)?)?.squeeze(D::Minus1);
let shape = scores.shape();
causal.broadcast_as(shape)?.where_cond(&scores, &ninf.broadcast_as(shape)?);

scores.broadcast_add(mask)         F32 [1, 2, 3, 3] [0.0, -inf, -inf, 3.0, 4.0, -inf, 6.0, 7.0, 8.0, 9.0, -inf, -inf, 12.0, 13.0, -inf, 15.0, 16.0, 17.0]
x.matmul(w.t()).broadcast_add(bias) F32 [2, 4] [3.5, 12.5, 21.5, 30.5, 3.5, 12.5, 21.5, 30.5]
w.matmul(v.unsqueeze(1)).squeeze(1) F32 [4] [3.0, 12.0, 21.0, 30.0]
cond, -inf broadcast_as(scores)    F32 [1, 2, 3, 3] [0.0, -inf, -inf, 3.0, 4.0, -inf, 6.0, 7.0, 8.0, 9.0, -inf, -inf, 12.0, 13.0, -inf, 15.0, 16.0, 17.0]

和 PyTorch 的值逐个相同（第一行的前 9 个就是 PyTorch 打印的第 0 个头）。几点：

• broadcast_as 得到的是步长为 0 的视图，不复制，可以直接交给 where_cond。
• 加偏置不用自己写：candle_nn::Linear 的 forward 最后就是 x.broadcast_add(bias)。
• [b, s, k] @ [k, n] 这种秩不同的乘法用 broadcast_matmul，或者先把批维并进行（Linear 的做法），见形状与广播的对应表。批维用 expand 扩出来再 matmul 有一个上游 bug，见expand 之后 matmul，后几批算错。
• masked_fill 在 candle 里没有，where_cond 加 broadcast_as 就是它的写法；candle-transformers 的 llama.rs 自带一个这样的 masked_fill 小函数，见归约与索引的「where_cond：顺序相同，不广播」。

怎么检查

• 读到 shape mismatch in add（或 sub、mul、div），先看两边的秩是否不同、或者有一维是 1：是的话就是该换 broadcast_*。报 shape mismatch in broadcast_add 才是形状真的不兼容。
• 从 PyTorch 搬代码时，凡是 @ 的一边是一维的，都要写 unsqueeze；torch.where 的三个参数里有标量或低秩张量的，都要 broadcast_as。
• 这类错误都是 Err，第一次运行就会暴露，用一组小形状（批 1、头 2、长度 3）跑一遍前向就能全部找出来。

小结

• candle 的 +、* 要求形状完全相同，matmul 不收一维，where_cond 三个形状必须相同；PyTorch 在这三处都隐式广播或补维。
• 报错依次是 shape mismatch in add、in matmul、in where_cond，都是 Err，CPU 和 Metal 相同。
• 改法：broadcast_add 等，unsqueeze(1) 再 squeeze，条件和标量 broadcast_as 到目标形状。