einsum 的梯度形状不对
主线位置：einsum 线的最后一步（AD 规则），在 eager 主线的
backward()里触发。eager 的 einsum 什么时候走带 AD 规则的扩展算子、规则怎么注册，见 semantic AD 与扩展；规则的代码逐行分析见 einsum 的「长度 1 广播时梯度形状错」。
flowchart TB a["a 形状 2 1"] --> fwd["前向 ij,jk->ik"] b["b 形状 3 2"] --> fwd fwd --> bc["j 从 1 广播到 3"] bc --> vjp["反向 ik,jk->ij"] vjp --> g["梯度形状 2 3"] g --> miss["缺少沿 j 求和"] fix["先 broadcast_in_dim"] --> okg["梯度形状 2 1"]
实验 crate exp-tenferro8/pit1-pit，bin p12_bcast_grad，release；CpuEngine::serial()；M2 Max、macOS 27、rustc 1.99.0。
现象
a 的形状是 [2, 1]，b 是 [3, 2]。"ij,jk->ik" 里 j 在 a 中长 1、在 b 中长 3，前向按 NumPy 的规则把 a 沿 j 广播。对 a 求梯度：
let ad = AdContext::builder()
.with_semantic_extension_rules(tenferro_einsum::semantic_ad_rules()?)?
.build()?;
let rt = EagerRuntime::with_ad_context(&ad)?;
let mut engine = CpuEngine::serial();
let a = rt.variable_from(Tensor::from_vec_col_major(vec![2, 1], vec![1.0_f64, 2.0])?)?;
let b = rt.constant_from(Tensor::from_vec_col_major(vec![3, 2], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0])?)?;
let loss = rt.with_eager_session(&mut engine, |s| {
let out = s.einsum(&[&a, &b], "ij,jk->ik")?;
Ok::<_, tenferro_einsum::Error>(s.reduce_sum(&out, None)?)
})?;
loss.backward(&mut engine)?;
eager ij,jk->ik a[2, 1] b[3, 2]: loss=[63.0]
a.grad shape=[2, 3] value=[5.0, 5.0, 7.0, 7.0, 9.0, 9.0]
finite difference = [21.000000000270802, 21.000000000270802]
eager ...j,jk->...k a[2, 1] b[3, 2]: loss=[63.0]
a.grad shape=[2, 3] value=[5.0, 5.0, 7.0, 7.0, 9.0, 9.0]
finite difference = [21.000000000270802, 21.000000000270802]
梯度的形状和 a 不一样，backward() 没有报错。有限差分（中心差分，步长 1e-6，在具体张量上算 einsum 再求和）给出 [21, 21]，正好是错误梯度沿 j 的和 5 + 7 + 9。写成省略号也一样。
梯度接下来去哪，错误就带到哪。手写一步梯度下降 a - 0.1 * grad（在 engine 上用直接运算）：
a - 0.1 * grad: shape [2, 3]
check grad shape == a shape: false
逐元素运算按 NumPy 规则再广播一次，参数悄悄从 [2, 1] 长成 [2, 3]，仍然没有报错。
出处
eager 的 einsum_subscripts_with_broadcast 先试一次 dot_general 捷径，再在 requires_broadcast 为假时把收缩展开成核心算子，这两条路求导都用核心算子的规则。requires_broadcast 发现同一个标签一边长 1、另一边不是 1 时，才建 einsum 扩展算子，求导用 einsum 自己的规则。所以只有带长度 1 广播的输入会踩到这个坑，写不写省略号都一样。
反向规则在 semantic_einsum_vjp。对第 active_idx 个输入，余切是一个新的 einsum，输出标签是这个输入的标签里「在输出或其他输入里出现过」的那些：
let vjp_output_labels: Vec<u32> = input_labels[active_idx]
.iter()
.copied()
.filter(|label| available_labels.contains(label))
.collect();
只有这组标签和输入自己的标签不同时，才调 semantic_broadcast_einsum_vjp 补一步。以 a 为例，j 也出现在 b 里，输出标签就是 [i, j]，和 a 的标签相同，什么都不补；反向 einsum "ik,jk->ij" 的 j 从 b 那里取到长度 3，结果就是 [2, 3]。规则只比较标签，不比较长度，前向允许的「同一标签长度 1 对长度 n」在反向没有对应的「沿这个轴求和、保持维数」。
这是上游 bug，0.8.0 上复现，结果和 0.7.1 逐位相同。广播随 PR #1705（v0.4.0）的省略号一起加进来，当时的反向测试没有覆盖长度 1（einsum 的冷知识）。三个输入时同样的广播在 0.8.0 前向就失败了，见 label 106 has extents 1 and 3。traced 入口不写省略号时在编译期就拒绝这种广播。
修法
在 einsum 之前把广播做出来，让广播自己的反向规则去求和。两种写法，同一个程序：
// 一：先显式广播到 [2, 3]
let loss = rt.with_eager_session(&mut engine, |s| {
let wide = s.broadcast_in_dim(&a1, &[2, 3], &[0, 1])?;
let out = s.einsum(&[&wide, &b], "ij,jk->ik")?;
Ok::<_, tenferro_einsum::Error>(s.reduce_sum(&out, None)?)
})?;
// 二：去掉长度 1 的轴，下标照实写
let loss = rt.with_eager_session(&mut engine, |s| {
let flat = s.reshape(&a2, vec![2])?;
let out = s.einsum(&[&flat, &b], "i,jk->ik")?;
Ok::<_, tenferro_einsum::Error>(s.reduce_sum(&out, None)?)
})?;
broadcast_in_dim first: a.grad shape [2, 1] [21.0, 21.0]
reshape + "i,jk->ik": a.grad shape [2, 1] [21.0, 21.0]
两种都和有限差分一致。第二种更清楚：a 本来就和 j 无关，"i,jk->ik" 把这件事写在下标里，也省掉一次物化广播；两个输入都不再需要广播，einsum 不走扩展算子，求导用核心算子的规则。第一种适合长度 1 的轴是「有时 1、有时 n」的批维，broadcast_in_dim 的 VJP 会沿被广播的轴求和并保持形状。
上游修好之前，凡是要对它求导的 einsum 输入，都不要依赖隐式的长度 1 广播，包括省略号的批维。（闭包里的 Ok::<_, tenferro_einsum::Error> 是必需的，原因见 From<tenferro_einsum::Error> is not implemented。）
怎么确认没中招
backward()之后断言每个变量的梯度形状等于它自己的形状：assert_eq!(a.grad()?.unwrap().shape(), a.shape())。这一行就能抓住这个 bug。- 对小尺寸输入做一次有限差分核对（上面的
fd只有十几行），尤其是带长度 1 的维度时。 - 搜 einsum 的调用点，看有没有输入的某一维是 1 而另一个输入同一标签的维度大于 1。
小结
- einsum 前向接受长度 1 的广播，反向规则只看标签不看长度，不在被广播的轴上求和：梯度形状变成广播后的形状，值也错，
backward()不报错。 - 只有需要长度 1 广播、因而走 einsum 扩展算子的调用受影响；写省略号也一样。
- 上游 bug（0.8.0 复现）。先
broadcast_in_dim，或reshape掉长度 1 的轴并改写下标；测试里断言梯度形状。