Rust in Everything
基于 tenferro-linalg 0.8.0

SVD 的输出

主线位置：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 --> err["反向报错"]
  meta --> rs["reshape 成 4"]
  rs --> sl2["slice"]
  sl2 --> ok["反向正常"]
  meta --> tk["take_axis 或掩码"]
  tk --> ok
  ok --> sq["sqrt s"]
  sq --> zero["s 恰为 0 时 NaN"]

TRG 的一步在 JAX 里写成：

ua, sa, vah = svd(ma)
s1 = (ua[:, :k] * jnp.sqrt(sa[:k])).reshape(d, d, k)
s3 = (vah[:k, :].T * jnp.sqrt(sa[:k])).reshape(d, d, k)

搬到 tenferro，这三行会遇到三件事：SVD 默认的形状不同；直接对 SVD 的输出切片，前向没问题，反向报错；保留下来的奇异值里如果有一个恰好是 0，sqrt 的梯度是 NaN。

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

薄的还是满的

jnp.linalg.svd(a) 默认 full_matrices=True，6×4 的矩阵得到 6×6 的 U；这样的 SVD 在非方阵上不能求导。tenferro 的 svd() 是薄的，满的是另一个函数 svd_full()，同样不能求导：

# JAX
full_matrices=True shapes: [(6, 6), (4,), (4, 4)]
grad through jnp.linalg.svd(a) (full_matrices=True): NotImplementedError: Singular value decomposition JVP not implemented for full matrices
# tenferro
svd_full shapes u Some([6, 6]) s Some([4]) vt Some([4, 4])
svd (thin) shapes u Some([6, 4]) s Some([4]) vt Some([4, 4])
grad through svd_full u: Err(unsupported vjp AD rule for tenferro-linalg.linalg.v1)

所以 JAX 里要对非方阵的 SVD 求导，本来就得写 full_matrices=False，它对应 tenferro 的 svd()。

切片之后求不了导

同一个 4×4 矩阵，对 sum(s[:2]) 求导。JAX 直接写：

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

grad of sum(s[:2]), column 0: [-0.4013948101174394, 0.17984649334769795, -0.3375913943613526, 0.7146239408315301]

tenferro 一样写，slice 本身成功，形状也是 [2]，但 ad.grad 报错：

let (u, s, vt) = x.svd()?;
let l = s.slice(SliceConfig { starts: vec![0], limits: vec![2], strides: vec![1] })?.reduce_sum(None)?;
ad.grad(&l, &x)

sum(s[:2])                         grad Err(vjp (GraphBuild): runtime state failure: semantic Vjp does not support derivative metadata: slice input requires constant exact or upper-bound extents)
sum(u[:, :2])                      grad Err(vjp (GraphBuild): runtime state failure: semantic Vjp does not support derivative metadata: slice input requires constant exact or upper-bound extents)
sum(vt[:2, :])                     grad Err(vjp (GraphBuild): runtime state failure: semantic Vjp does not support derivative metadata: slice input requires constant exact or upper-bound extents)
sum(q[:, :2])                      grad Err(vjp (GraphBuild): runtime state failure: semantic Vjp does not support derivative metadata: slice input requires constant exact or upper-bound extents)
sum(r[:2, :])                      grad Err(vjp (GraphBuild): runtime state failure: semantic Vjp does not support derivative metadata: slice input requires constant exact or upper-bound extents)

U、s、Vᵀ 和 QR 的 Q、R 都一样。eager 下 backward 报同一个错（semantic_eager_vjp (GraphBuild): ...）。只做前向、或者用 ad.jvp 正向求导都没事：

jvp of sum(s[:2]): Ok [-0.8939366745326542]

JAX 的 jax.jvp 给出 −0.8939366745326536，两边差 6e-16。

原因：长度是 Min(4, 4)

打印形状元数据：

x       concrete Some([4, 4]) sym Some([SymDim(Const(4)), SymDim(Const(4))])
svd u   concrete Some([4, 4]) sym Some([SymDim(Const(4)), SymDim(Min(Const(4), Const(4)))])
svd s   concrete Some([4]) sym Some([SymDim(Min(Const(4), Const(4)))])
svd vt  concrete Some([4, 4]) sym Some([SymDim(Min(Const(4), Const(4))), SymDim(Const(4))])
qr q    concrete Some([4, 4]) sym Some([SymDim(Const(4)), SymDim(Min(Const(4), Const(4)))])
qr r    concrete Some([4, 4]) sym Some([SymDim(Min(Const(4), Const(4))), SymDim(Const(4))])
...
s.reshape([4]) sym Some([SymDim(Const(4))])

try_concrete_shape 能算出 4，但图里记的是一个表达式 Min(Const(4), Const(4))。它来自 svd_meta，薄 SVD 的 k = min(m, n) 按符号写，两边都是常数时也不化简：

fn svd_meta(dtype: DType, shape: &[SymDim]) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>> {
    let (m, n, batch) = matrix_meta_parts("tenferro-linalg.svd", shape)?;
    let k = m.clone().min(n.clone());
    Ok(vec![
        (dtype, matrix_shape(m, k.clone(), batch)),
        (singular_values_dtype(dtype), vector_shape(k.clone(), batch)),
        (dtype, matrix_shape(k, n, batch)),
    ])
}

qr_meta 也是这样。slice 的反向要把余切补零补回输入的形状，slice_vjp 先用 const_usize_shape_for_inverse_slice_padding 读输入的长度，只认字面常数：

        .map(|extent| match extent {
            ShapeExtent::Exact(DimExpr::Const(value))
            | ShapeExtent::UpperBound(DimExpr::Const(value)) => Ok(*value),
            ShapeExtent::Exact(_) | ShapeExtent::UpperBound(_) | ShapeExtent::Unknown => {
                Err(metadata_error(
                    SemanticTransformRole::Vjp,
                    format!("{field} requires constant exact or upper-bound extents"),
                ))
            }
        })

Min(Const(4), Const(4)) 不是 Const，于是报错。这是元数据没有常量折叠造成的，算一个上游的缺口；本卷没有报上去。

❗IMPORTANT
动脑筋：eigh 的输出切片后再求导，比如 sum(v[:, 2:]) 和 sum(w[2:])，会不会报同样的错？

提示：eigh 的输入是方阵，看看它的形状元数据需不需要 min。

三种绕法

// 1. 先 reshape 成同样的具体形状，长度变成 Const(4)
let s4 = s.reshape([4])?;
let l1 = s4.slice(SliceConfig { starts: vec![0], limits: vec![2], strides: vec![1] })?.reduce_sum(None)?;
// 2. 乘一个 0/1 掩码，不改变形状
let mask = TracedTensor::from_vec_col_major(vec![4], vec![1.0_f64, 1.0, 0.0, 0.0])?;
let l2 = (&s * &mask)?.reduce_sum(None)?;
// 3. take_axis：按给定下标取（底层是 gather）
let l3 = s.take_axis(0, &[0, 1])?.reduce_sum(None)?;

sum(s.take_axis(0, &[0, 1]))       grad Ok
sum(u.take_axis(1, &[0, 1]))       grad Ok
sum(s.reshape([4])[:2])            grad Ok
...
reshape-then-slice grad[..4] [-0.40139481011743905, 0.1798464933476977, -0.3375913943613522, 0.71462394083153]
mask grad[..4]               [-0.40139481011743905, 0.1798464933476977, -0.3375913943613522, 0.71462394083153]
max |reshape - mask| = 0e0
take_axis grad: max |take - mask| = 0e0
...
eager reshape-then-slice backward Ok; max |eager - traced| = 0e0

三种写法的梯度逐位相同，第一列和 JAX 的 s[:2] 梯度相差不超过 4e-16。eager 下 reshape 再切片也能 backward，和 traced 逐位相同。

  • reshape 再切片最接近 JAX 的写法，试验阶段的 TRG 代码用的就是它：u.reshape([m, m])、s.reshape([m])、vt.reshape([m, m]) 之后再取前 k 列。
  • 掩码不改变形状，后面的收缩还是 m 维，多算了被截掉的部分。tenferro 卷用它绕开 动态截断的梯度问题。
  • take_axis 不用先知道是不是 Min，但它是 gather，反向是 scatter，比连续切片多一份下标。

上游指南推荐的写法 slice_axis 底下还是 slice，同样报错：

sum(s.slice_axis(0, 0..2))         grad Err(vjp (GraphBuild): runtime state failure: semantic Vjp does not support derivative metadata: slice input requires constant exact or upper-bound extents)
sum(u.slice_axis(1, 0..2))         grad Err(vjp (GraphBuild): runtime state failure: semantic Vjp does not support derivative metadata: slice input requires constant exact or upper-bound extents)
📝NOTE
冷知识：截断的「官方写法」是怎么来的

PR #1338（「Add linalg AD route and decomposition APIs」，首次发布于 v0.3.0）之前，traced 张量上取一段只有 slice(SliceConfig { starts, limits, strides })。这个 PR 处理了要求按轴切片的 issue #1336，随附的工作记录写着：「Added traced and eager slice_axis, slice_builder, and builder range/step/take operations. Added traced take_axis as a usize wrapper over index_select.」以及「Documented fixed-rank truncation as svd() plus slice_axis/take_axis」。同一天的 #1339 又明确「No SVD truncation policy; users continue to slice returned factors.」于是 0.8.0 的线性代数指南里，截断 SVD 的示例就是 u.slice_axis(1, 0..rank) 这样的三行，只断言形状。上面的实验说明，这三行在 0.8.0 里不能求导，同一句话里的 take_axis 可以。这条路在 v0.3.0 时能不能求导，本卷没有验证。

k 由数据决定

按阈值截断（保留 s > tol 的那些）时，k 是数据算出来的。JAX 在 jit 里不接受动态长度：

def trunc(a, tol):
  u, s, vh = jnp.linalg.svd(a, full_matrices=False)
  k = jnp.sum(s > tol)
  return u[:, :k] * s[:k] @ vh[:k]

jax.jit(trunc)(x, 0.5)

IndexError: Array slice indices must have static start/stop/step to be used with NumPy indexing syntax. Got slice(None, JitTracer(int64[]), None) at position 1. To index an array at a dynamic position with a static slice size, use x[jax.ds(start, size)] or lax.dynamic_slice/dynamic_update_slice instead (JAX does not support dynamically sized arrays within traced functions).

不 jit 时这段能跑。jit 里的常见做法是掩码：保持 m 维，把 s ≤ tol 的那些乘 0。

tenferro 有 JAX 没有的东西：s.dynamic_truncate(&k, 0) 的输出长度是一个上界，执行时才确定，一个程序能处理不同的 k。但对二维以上的截断结果求导会在执行时失败，tenferro 卷的 线性代数「截断 SVD 与动态形状」一节实测过，常见坑 reshape element-count mismatch 有最小复现。要对按阈值截断的 SVD 求导，两边都用掩码。

TRG 主线没有这个问题：每一步的 k 是 [2, 4, 4, 16, …]，事先定好，和 β 无关（试验阶段从奇异值的分布里看出来的结构秩）。

恰好为 0 的奇异值和 sqrt

TRG 第一步的矩阵 4×4、秩是 2。如果按 k = min(D², χ) 的「朴素」截断，第一步保留全部 4 个奇异值，再开平方根。tenferro 在这里给 NaN，JAX 不给：

# tenferro (faer)
TRG step-0 matrix, beta=0.3: faer s = ["2.3709304364845347e0", "1.2733071642964824e0", "5.329577449131604e-17", "-0e0"]
  bits of s[3] = 0x8000000000000000, is_sign_negative = true, s[3] == 0.0: true
# JAX (LAPACK)
TRG step-0 matrix, beta=0.3: LAPACK s = ['2.370930e+00', '1.273307e+00', '4.106732e-17', '3.505205e-17']

faer 把最后一个奇异值算成了恰好的零，而且是 −0（符号位是 1）。LAPACK 给的是 1e-17 量级的正数。对这些奇异值的平方根求和，再对 β 求导：

# tenferro
all 4, sqrt                        value [2.668191731038373] grad [NaN]
all 4, sqrt, derivative_eps 1e-6   value [2.668191731038373] grad [NaN]
first 2, sqrt                      value [2.6681917237379777] grad [2.9280661736460765]
all 4, safe sqrt                   value [2.6681917237379777] grad [2.9280661736460765]
all 4, single outer where          value [2.6681917237379777] grad [NaN]
# JAX，SVD 用 custom_jvp 的展宽规则
  all 4, sqrt        value 2.6681917360668335 grad 2.9280661918647124
  first 2, sqrt      value 2.6681917237379777 grad 2.9280661736460765
  all 4, safe sqrt   value 2.6681917237379777 grad 2.9280661736460765

  • tenferro 的 NaN 出在 sqrt 上，不在 SVD 的规则里：d√s/ds = 1/(2√s) 在 s = 0 时是无穷，这个无穷沿反向传回去，变成 NaN。只取前两个时梯度正常，说明 SVD 本身的反向没问题。改 derivative_eps 没用，它只管 SVD 规则里的 F。

  • 和 −0 的符号无关。JAX 在一个精确为 0 的奇异值上也一样（diag(2, 1, 0, 0) 沿一个方向求导）：

    diag(2,1,0,0): LAPACK s = ['2.000000e+00', '1.000000e+00', '0.000000e+00', '0.000000e+00']
  lorentz, sqrt      d/dt at t=0: nan
  lorentz, safe sqrt d/dt at t=0: 0.6137752937474935

  • JAX 在 TRG 上没出 NaN，只是因为 LAPACK 算出的是 4e-17 而不是 0。这时 1/(2√s) 约 8e7，梯度有限，但和只取前两个时差 1.8e-8：这两个奇异值是舍入噪声，它们的导数也是噪声。

  • 这和常见坑 SVD 的梯度全是 NaN 不是一回事：那里是非方阵的投影修正项除以 s，这里的矩阵是方的，SVD 本身的梯度没问题。

两种修法，两边通用：

  • 按秩截断：不保留结构上为 0 的奇异值。上面 first 2 就是，本卷 TRG 的 k = [2, 4, 4, 16, …] 也是这么来的。
  • 安全的 sqrt（JAX 里常说的 double where）：被丢掉的位置先换成 1 再开方，最后再置 0，反向里就不会出现 1/√0。
def safe_sqrt(s, tol=1e-14):
  keep = s > tol
  return jnp.where(keep, jnp.sqrt(jnp.where(keep, s, 1.0)), 0.0)

fn safe_sqrt(s: &TracedTensor, n: usize, tol: f64) -> R<TracedTensor> {
  let tolv = TracedTensor::from_vec_col_major(vec![n], vec![tol; n])?;
  let one = TracedTensor::from_vec_col_major(vec![n], vec![1.0_f64; n])?;
  let zero = TracedTensor::from_vec_col_major(vec![n], vec![0.0_f64; n])?;
  let keep = s.compare(&tolv, CompareDir::Gt)?;
  let inner = TracedTensor::where_select(&keep, s, &one)?;
  Ok(TracedTensor::where_select(&keep, &inner.sqrt()?, &zero)?)
}

只写外层的 where 不够，上面 tenferro 的 single outer where 一行就是 NaN，JAX 同样写法在 diag(2, 1, 0, 0) 上也是 nan：sqrt(0) 照样在图里，反向时 where 给它的余切是 0，0 · ∞ 还是 NaN。

# JAX，jnp.where(s > 1e-14, jnp.sqrt(s), 0.0)，diag(2, 1, 0, 0)
single where d/dt: nan

放到整个 TRG 的一步里（d/dβ [Σ T′ + ln‖T‖]），朴素截断加安全的 sqrt 和按秩截断给出同一个梯度：

k=2 eps=None safe_sqrt=false: value [2.8236462733691012] grad [2.8873891988539446]
k=3 eps=None safe_sqrt=false: value [2.8236462733691012] grad [2.8873891988539473]
k=4 eps=None safe_sqrt=false: value [2.8236462733691012] grad [NaN]
k=4 eps=Some(1e-6) safe_sqrt=false: value [2.8236462733691012] grad [NaN]
k=4 eps=None safe_sqrt=true : value [2.8236462733691017] grad [2.8873891988539446]

k = 3 留下了 5.3e-17 那个，梯度末几位变了（噪声），但还是有限的；k = 4 留下了 −0，就是 NaN。

安全的 sqrt 等于把阈值以下的奇异值当作结构上的 0。它们真的会随参数离开 0 时（比如 diag(2, 1, 0, 0) 加一个任意方向的扰动），Σ√sᵢ 在这一点本来就不可导，上面的 0.6138 只是保留下来那两项的导数（w₀₀/(2√2) + w₁₁/2 = 0.6137752937474935，逐位相同）。

动脑筋答案

sum(v[:, 2:])  (eigh)              grad Ok
sum(w[2:])     (eigh)              grad Ok
...
eigh w  concrete Some([4]) sym Some([SymDim(Const(4))])
eigh v  concrete Some([4, 4]) sym Some([SymDim(Const(4)), SymDim(Const(4))])
  • 不会报错。eigh 的特征值、特征向量的长度都是输入的边长 n，元数据直接写 Const(4)，没有 Min。
  • 所以问题不在「分解的输出」，只在长度写成 min(m, n) 的那些：SVD 的 U、s、Vᵀ，QR 的 Q、R。

小结

  • jnp.linalg.svd 默认满矩阵，求导要 full_matrices=False；tenferro 的 svd() 是薄的，svd_full() 不能求导。
  • SVD、QR 的输出长度记作 Min(Const(m), Const(n))，slice 的反向只认 Const，直接切片再求导报 slice input requires constant exact or upper-bound extents（traced 和 eager 都是；slice_axis 也是；前向和 jvp 没事）。先 reshape 成同样的具体形状、乘掩码或用 take_axis，三种梯度逐位相同。eigh 的输出不受影响。
  • k 由数据决定时，JAX 的 jit 不接受动态切片；tenferro 有 dynamic_truncate，但二维以上的截断结果不能求导。两边都用掩码。
  • faer 会给出恰好为 −0 的奇异值（LAPACK 给 1e-17 量级），保留它再开方，梯度是 NaN，derivative_eps 管不到。按秩截断，或者写双层 where 的安全 sqrt。

💬 关联讨论

发起讨论
加载中...