JAX 自带 SVD 的梯度从第 5 步起是 NaN

主线位置：TRG 主线「分段求导得内能」：搭 JAX 参考实现，用它的梯度去核对 tenferro。

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

flowchart TB
  ref["JAX 参考 TRG"] --> bi["jnp.linalg.svd"]
  bi --> s4["4 步以内 正确"]
  bi --> s5["第 5 步起 NaN"]
  s5 --> why["228 个奇异值相等"]
  ref --> cj["custom_jvp 展宽"]
  cj --> sq["平方差 不屏蔽"]
  sq --> fma["jit 下 1e14 到 1e26"]
  cj --> pr["写成乘积 加 where"]
  pr --> ok["和 tenferro 一致"]

这个坑在 JAX 一侧：要拿 JAX 的梯度去核对 tenferro，先得让 JAX 的梯度是对的。实验代码：JAX 一侧是 p20_builtin_svd.py（trg_jax.py 里的 TRG 和 lorentz_svd），tenferro 一侧是动手页的 eager 分段程序 ising_trg（release，CpuEngine::serial()）。β = 0.4，χ = 16，截断数 k = [2, 4, 4, 16, …]。M2 Max、macOS 27、rustc 1.99.0。

现象

参考实现里的 SVD 直接用 jnp.linalg.svd，对 lnZ/N 求 jit(grad)，和对值的中心差分（h = 1e-5）比：

ua, sa, vah = jnp.linalg.svd(ma, full_matrices=False)
...
jax.jit(jax.grad(make_trg(ks, SVD["builtin"])))(b)

steps=3 fd=1.499230514157 builtin=1.499230514558 ...
steps=4 fd=1.379116481881 builtin=1.379116482259 ...
steps=5 fd=1.295053391531 builtin=nan ...
steps=6 fd=1.222308598453 builtin=nan ...
step 4 matrix (256, 256): most repeated singular value 7.124e-17 x 228

同样的 5 步，tenferro 默认的 svd()：

tenferro 0.8.0 eager staged chi=16 steps=5 width=1 h=1e-5 ks=[2, 4, 4, 16, 16]
beta=0.4000 lnZ/N=0.892077682813751 u=-1.295053391806669 c=0.987162872869 | value_and_grad 40.1 ms

u = −∂(lnZ/N)/∂β，所以 tenferro 的梯度是 1.295053391807，和差分差 2.8e-10（差分的截断误差）。两边一比，JAX 是 NaN、tenferro 是有限值，很容易先怀疑 tenferro。

出处

第 5 步的矩阵是 256×256，有 228 个奇异值在 7e-17 处逐位相等（它们本来都是 0，舍入后落在同一个数上）。JAX 内置的 SVD 反向规则里有 1/(s_j² − s_i²)，相等的奇异值让分母恰好为 0，0·∞ 得到 NaN 并扩散到整个梯度。tenferro 默认的规则把它换成洛伦兹展宽 F = d/(d² + ε²)，ε = 1e-12，d = 0 时 F = 0。为什么是第 5 步、最小的 4×4 复现、PyTorch 的同样结果和展宽规则的推导，见 简并奇异值。

tenferro 卷的 SVD 的梯度全是 NaN 是另一个机制：非方阵恰好有一个零奇异值时，m ≠ n 的投影修正项除以 s。TRG 的矩阵都是方阵，不走那一项。

修法

在 JAX 里用 jax.custom_jvp 把 SVD 的规则换成和 tenferro 相同的展宽形式（trg_jax.py 的 lorentz_svd）。照教科书写 d = s_j² − s_i² 还会踩第二个坑：

d = s[None, :] ** 2 - s[:, None] ** 2               # 教科书写法
f = d / (d * d + EPS * EPS)

d = (s[None, :] - s[:, None]) * (s[None, :] + s[:, None])   # 改成乘积
f = jnp.where(jnp.eye(s.shape[0]) > 0, 0.0, d / (d * d + EPS * EPS))

steps=3 fd=1.499230514157 builtin=1.499230514558 custom_jvp=1.499230514558 custom_jvp(s**2, jit)=3.777812e+14
steps=4 fd=1.379116481881 builtin=1.379116482259 custom_jvp=1.379116482259 custom_jvp(s**2, jit)=3.982413e+14
steps=5 fd=1.295053391531 builtin=nan custom_jvp=1.295053391807 custom_jvp(s**2, jit)=-3.050596e+20
steps=6 fd=1.222308598453 builtin=nan custom_jvp=1.222308598298 custom_jvp(s**2, jit)=5.144028e+26

• 乘积写法加 where 屏蔽对角线（custom_jvp 列）：每一步都和差分对得上，5 步时是 1.295053391807，和 tenferro 打印的 u 只差末位的舍入（-1.295053391806669）。
• 平方差写法、对角线不屏蔽（最后一列）在 jit 下从第 3 步起就是 1e14 到 1e26：XLA 把 s·s − s·s 编成一条融合乘加，对角线不是 0 而是 1e-16 量级，除以 ε² 后放大。没有 NaN，比 NaN 更难发现。机制见 简并奇异值 的「jit 的 FMA 陷阱」。

怎么确认没中招

• 参考实现的梯度先和它自己的值的中心差分比一次，每加一步比一次，再拿它去核对 tenferro。
• 看第一次出 NaN 的那一步的奇异值：有成片逐位相等的值，就是内置规则的问题，不是 tenferro 的问题。
• 自定义规则在 jit 和不 jit 下各跑一次：结果差几个数量级，就是 FMA。

小结

• JAX 内置的 SVD 规则遇到逐位相等的奇异值给 NaN；TRG 第 5 步有 228 个。tenferro 默认的展宽规则给出正确梯度。
• JAX 参考实现用 custom_jvp 换成同一个展宽规则；d 写成 (s_j − s_i)(s_j + s_i)，对角线用 where 置 0。
• 平方差写法又不屏蔽对角线时，jit 下不报错、不出 NaN，梯度到 1e14 以上。