sqrt(s) 的梯度是 NaN

主线位置：TRG 主线「一步 SVD 截断加收缩」：sqrt(s[:k]) 把奇异值分到两边。

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

flowchart TB
  m["秩 2 的 4x4 矩阵"] --> lp["LAPACK"]
  lp --> tiny["s4 约 3.5e-17"]
  tiny --> fin["sqrt 梯度有限"]
  m --> fa["faer"]
  fa --> z["s4 = -0"]
  z --> nan["sqrt 梯度 NaN"]
  z --> safe["双 where 的 sqrt"]
  safe --> ok["梯度正确"]

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

现象

TRG 第一步要把 2⁴ 的初始张量重排成 4×4 矩阵做 SVD。截断数按 k = min(D², χ) 取，χ ≥ 4 时第一步保留全部 4 个奇异值，再开方。β = 0.3，对 Σ sqrt(s) 求 d/dβ。JAX：

step0 = lambda b: initial_tensor(b).transpose(2, 1, 0, 3).reshape(4, 4)
f = lambda b: jnp.sum(jnp.sqrt(svd(step0(b))[1]))
jax.jit(jax.grad(f))(b)

LAPACK s = ['2.370930e+00', '1.273307e+00', '4.106732e-17', '3.505205e-17']
builtin    sum(sqrt(s)) value 2.6681917360668335 d/dbeta 2.9280661918647124
custom_jvp sum(sqrt(s)) value 2.6681917360668335 d/dbeta 2.9280661918647124
first 2 only       value 2.6681917237379777 d/dbeta 2.9280661736460765

tenferro 照搬：

let ma = initial_tensor(&beta)?.transpose(&[2, 1, 0, 3])?.reshape([4, 4])?;
let s = ma.svd()?.1.reshape([4])?;
let l = s.sqrt()?.reduce_sum(None)?;
ad.grad(&l, &beta)

faer s = ["2.3709304364845347e0", "1.2733071642964824e0", "5.329577449131604e-17", "-0e0"]
s[3] bits 0x8000000000000000
sum(sqrt(s))                         value [2.668191731038373] d/dbeta [NaN]
sum(sqrt(s)), derivative_eps 1e-6    value [2.668191731038373] d/dbeta [NaN]

• 这个矩阵的秩是 2，后两个奇异值在数学上是 0。LAPACK 给出两个 1e-17 量级的正数；faer 给出一个 1e-17 和一个精确的 -0（符号位为 1）。
• JAX 的梯度有限；tenferro 是 NaN，加大 derivative_eps 也一样。

出处

sqrt 在 0 处的导数是无穷：反向里是 ds / (2·sqrt(s))，s = −0 时分母为 0。这和 SVD 的求导规则无关，derivative_eps 只管奇异值简并时的展宽（简并奇异值），碰不到 sqrt。JAX 一侧没出 NaN，只是因为 LAPACK 没有算出精确的 0；把矩阵换成 diag(2, 1, 0, 0)，LAPACK 也给出 0，JAX 的梯度同样是 NaN（SVD 的输出）。

JAX 的 2.928066191865 也不对：真值是只算前两个的 2.928066173646（后两个奇异值恒为 0，对 β 的导数是 0），差 1.8e-8，来自两个 1e-17 的「奇异值」开方后的导数。tenferro 卷的 SVD 的梯度全是 NaN 是另一个机制（非方阵的投影项除以 s），那里 sqrt 不是原因。

修法

用「两次 where」的 sqrt：小于阈值的位置先换成 1 再开方，最后再置 0，这样被屏蔽的位置从来不碰 sqrt(0)：

fn safe_sqrt(s: &TracedTensor, n: usize, tol: f64) -> R<TracedTensor> {
  let keep = s.compare(&full(n, tol)?, CompareDir::Gt)?;
  let inner = TracedTensor::where_select(&keep, s, &full(n, 1.0)?)?;
  Ok(TracedTensor::where_select(&keep, &inner.sqrt()?, &full(n, 0.0)?)?)
}

sum(sqrt(s + 1e-30))                 value [2.668191731038374] d/dbeta [2.92110713571551]
sum(safe_sqrt(s))                    value [2.6681917237379777] d/dbeta [2.9280661736460765]

（实验代码里 full 是三行 from_vec_col_major。）safe_sqrt 的结果和 JAX 只取前两个的值、梯度逐位相同。加一个小量的写法消掉了 NaN，梯度却差了 7e-3：1e-17 那个奇异值开方后的导数被放大了。从根上避开的办法是截断数按秩取：TRG 前三步的 k 取 2、4、4（本卷的 TRG 都这样做，见 Ising 模型的 TRG），零奇异值根本不进 sqrt。

怎么确认没中招

• 对每一步打印被保留的最小奇异值，看有没有 0e0 或 -0e0；faer 给出精确 0 的情况 LAPACK 往往给出 1e-17，所以 JAX 参考实现正常不说明 tenferro 正常。
• 梯度出 NaN 时先把 sqrt 换成 safe_sqrt 试一次：NaN 消失，就是这个坑，不是 SVD 规则的问题。

小结

• 秩亏的矩阵在 faer 里给出精确的 −0 奇异值，sqrt 的梯度在 0 处无穷，结果 NaN；derivative_eps 不管用。
• JAX 的 LAPACK 给出 1e-17 的小数，梯度有限但也偏了 1.8e-8。
• 截断数按秩取，或用两次 where 的 sqrt；不要加小量。