reshape element-count mismatch: from 6 to 12

主线位置：traced 主线的 grad 和执行，用在运行期才决定长度的程序上。dynamic_truncate 的形状元数据是上界（UpperBound），见 TracedTensor；截断 SVD 怎么用它、为什么上游例子用掩码，见 线性代数 的「截断 SVD 与动态形状」。

flowchart TB
  x["x 形状 4x3"] --> t["dynamic_truncate k=2"]
  t --> xk["x_k 实际 2x3"]
  xk --> m["x_k * x_k"]
  m --> g["grad 建图编译 通过"]
  g --> run["执行"]
  run --> err["reshape 6 对 12"]
  x --> mask["x 乘 0/1 掩码"]
  mask --> ok["梯度正确"]

实验 crate exp-tenferro8/pit2-exp，bin p25_dyn_trunc，CPU（CpuEngine::serial()），debug，M2 Max、macOS 27、rustc 1.99.0。

现象

x 是 4×3 的常数，k = 2 是图里的一个标量，保留前 k 行，对三种损失求导：

let x = TracedTensor::from_vec_col_major(vec![4, 3], (1..=12).map(|i| i as f64).collect())?;
let k = TracedTensor::from_vec_col_major(vec![], vec![2.0_f64])?;
let xk = x.dynamic_truncate(&k, 0)?;
for (name, l) in [
  ("sum(x_k)", xk.reduce_sum(None)?),
  ("sum(x_k * x_k)", (&xk * &xk)?.reduce_sum(None)?),
  ("sum(exp(x_k))", xk.exp()?.reduce_sum(None)?),
] {
  // 先算前向，再 l.grad(&x) 并执行
}

sum(x_k): loss = [33.0], grad: [1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0]
sum(x_k * x_k): loss = [247.0], grad: run Err(TypedTensorView::reshape_view: reshape element-count mismatch: from 6 to 12)
sum(exp(x_k)): loss = [30691.499012904802], grad: run Err(TypedTensorView::reshape_view: reshape element-count mismatch: from 6 to 12)
rank 1, sum(x_k * x_k): grad [2.0, 4.0, 0.0, 0.0, 0.0]

• 前向都对：前两行 1, 5, 9, 2, 6, 10 的平方和是 247。
• loss.grad(&x) 成功返回，梯度程序也能编译；执行时才失败。
• 报错的是 reshape_view，可代码里没有写任何 reshape。6 是截断后的元素数（2×3），12 是截断前的上界（4×3）。
• 一维输入（最后一行）和只求和的损失都正确；二维输入接上「反向要读截断值」的运算（乘、exp）就失败。按第 1 维截断、[1, 5] 这种只多一个长度为 1 的轴的输入也一样（线性代数 的 dyn_trunc_grad）。截断 SVD 正是这种用法：U、s、Vt 都截到 k，再乘回去。

出处

DynamicTruncate 本身的反向是用 PadToMatch 把余切补回输入的长度，这一步没问题。问题在截断值后面的运算：它们的反向按静态元数据把余切整理回自己输入的形状，而截断值的静态形状是上界 [4, 3]。专门处理动态轴的是 truncate_ad_value_to_dynamic_axes，ReduceSum 的 VJP 调了它，把广播出来的 [4, 3] 截回实际长度；乘法和 exp 的反向里没有对应的一步，执行时某个按上界形状做的 reshape 拿到了只有 6 个元素的实际余切。具体是哪一条反向规则发出的这个 reshape，本页没有逐条定位。

和 0.7.1 比，失败的条件没变，报错从 reshape: incompatible shapes: lhs=[2, 3], rhs=[4, 3] 换成了按元素数报的 TypedTensorView::reshape_view，少了两个形状。这是 0.8.0 仍在的上游 bug。

修法

不改变形状，改用 0/1 掩码。掩码由行号和 k 比较得到，k 仍然是图里的值，程序仍然只编译一次：

let rows = TracedTensor::from_vec_col_major(vec![4], vec![0.0_f64, 1.0, 2.0, 3.0])?;
let mask = rows.compare(&k, CompareDir::Lt)?.convert(DType::F64)?.broadcast_in_dim(&[4, 3], &[0])?;
let xm = (&x * &mask)?;
let loss = (&xm * &xm)?.reduce_sum(None)?;

mask: loss = [247.0]
mask: grad = [2.0, 4.0, 0.0, 0.0, 10.0, 12.0, 0.0, 0.0, 18.0, 20.0, 0.0, 0.0]

梯度是 2x，被截掉的两行是 0，正确。代价是被截掉的部分照样参与计算；截断 SVD 里 U、Vt 保持满尺寸，只把 s 乘掩码（上游例子 dynamic_shape_truncated_svd 就是这么写的），多算的是一次满尺寸的矩阵乘。

怎么确认没中招

• 搜 dynamic_truncate，看输入的秩和截断值后面接的运算：秩不小于 2、后面不是单纯的求和，就会中招。
• 只要 traced 程序里有动态形状，求导后一定实际执行一次梯度程序；grad 能返回、能编译都不说明什么。
• 用掩码写法时，和截断写法比一次前向的结果，确认掩码的方向（Lt 还是 Le）没写反。

小结

• 二维以上张量 dynamic_truncate 之后接乘法、exp 等运算，梯度执行时报 reshape_view: reshape element-count mismatch: from 6 to 12；一维、只求和时正常。0.8.0 仍在的上游 bug。
• 原因：ReduceSum 的反向会把余切截回实际长度，乘法、exp 这一类没有，梯度程序里留着按上界形状做的 reshape。
• 用比较得到的 0/1 掩码代替截断，梯度正确，程序仍只编译一次。