slice input requires constant exact or upper-bound extents

主线位置：TRG 主线「一步 SVD 截断加收缩」：从 U、s、Vᵀ 里取前 k 个。

版本：JAX 0.12.0.1（CPU，x64）、tenferro-rs 0.8.0（CPU，native，单线程）

flowchart TB
  svd["svd 输出"] --> meta["长度 Min 4 4"]
  meta --> sl["slice 前 k 个"]
  sl --> fwd["前向正常"]
  sl --> err["ad.grad 报错"]
  meta --> rs["reshape 成 4"]
  rs --> sl2["再 slice"]
  sl2 --> ok["梯度和 JAX 一致"]

实验代码：tenferro 一侧是 bin p21_svd_slice（release，CpuEngine::serial()），JAX 一侧是 p21_svd_slice.py。M2 Max、macOS 27、rustc 1.99.0。

现象

4×4 矩阵，损失 Σ sqrt(s[:2])，正是 TRG 一步里截断加开方的写法。JAX：

f = lambda x: jnp.sum(jnp.sqrt(jnp.linalg.svd(x, full_matrices=False)[1][:2]))
jax.grad(f)(x)

value 2.607050271637118
grad[:, 0] (column-major first 4) [-0.14795662  0.05925285 -0.12136266  0.29850654]
grad of sum(u[:, :2]) finite: True

照搬到 tenferro，[:2] 写成 slice：

let (u, s, _vt) = x.svd()?;
let l = s.slice(sl(&[0], &[2]))?.sqrt()?.reduce_sum(None)?;
eval(&rt, &mut eng, &l)?;   // 前向
ad.grad(&l, &x)             // 反向

s: concrete Some([4]) sym Some([SymDim(Min(Const(4), Const(4)))])
forward sum(sqrt(s[:2])) = [2.607050271637118]
grad: Err(vjp (GraphBuild): runtime state failure: semantic Vjp does not support derivative metadata: slice input requires constant exact or upper-bound extents)
grad of sum(u[:, :2]): Err(vjp (GraphBuild): runtime state failure: semantic Vjp does not support derivative metadata: slice input requires constant exact or upper-bound extents)

前向和 JAX 逐位相同；ad.grad 在建导数图时失败。u[:, :2] 也一样。

出处

s 的具体长度是 4，图里记的长度却是表达式 Min(Const(4), Const(4))：薄 SVD 的 k = min(m, n) 按符号写，两边都是常数也不化简。slice 的反向要把余切补零补回输入的长度，只认字面常数，于是在 core_structural.rs 里报这条错误。QR 的输出相同，eigh 的输出没有这个问题，tenferro 自己的线性代数指南里 slice_axis 截断的写法也中招；逐个输出的对照、元数据的来源（svd_meta）和这条 API 的来历见 SVD 的输出。

修法

切片之前 reshape 成同样的具体形状，把长度变成 Const(4)；或者用 take_axis 按下标取：

let s4 = s.reshape([4])?;
let l2 = s4.slice(sl(&[0], &[2]))?.sqrt()?.reduce_sum(None)?;
let g = ad.grad(&l2, &x)?;
let l3 = s.take_axis(0, &[0, 1])?.sqrt()?.reduce_sum(None)?;

s.reshape([4]): sym Some([SymDim(Const(4))])
reshape then slice: value [2.607050271637118] grad[0..4] [-0.14795661635587345, 0.05925285332680167, -0.12136266164526498, 0.2985065360536259]
take_axis: max|grad - reshape grad| = 0e0
max|jax - tenferro (reshape then slice)| = 1.804e-16, max|grad| = 0.299

和 JAX 差 1.8e-16（梯度最大 0.3，约 1 个 ulp）；take_axis 和 reshape 的结果逐位相同。本卷的 TRG 程序都是先 u.reshape([m, m])、s.reshape([m])、vt.reshape([m, m]) 再切。

怎么确认没中招

对分解的输出切片之前打印 sym_shape()：里面有 Min(...)，反向就过不去。
前向能算不代表能求导。写完一步就对它做一次 ad.grad，不要等整个程序搭完。
jvp 不受影响（SVD 的输出里 sum(s[:2]) 的 jvp 能算），所以只用前向模式验证过的代码，换成反向时还会撞上。

小结

SVD、QR 输出的长度在图里是 Min(Const, Const)，slice 的反向只认常数，报 slice input requires constant exact or upper-bound extents。
前向不报错，错误在 ad.grad 建导数图时出现。
先 reshape 成同样的形状再切，或用 take_axis；梯度和 JAX 差 1 个 ulp。

slice input requires constant exact or upper-bound extents

现象

出处

修法

怎么确认没中招

小结

💬 关联讨论

登录