Rust in Everything
基于 tenferro-runtime 0.8.0

线性代数

主线位置:直接计算主线上,a.svd(&mut engine) 和 matmul 一样借 engine 执行,只是落到的 provider 是 tlinalg(native)或 tlinalg-blas(blas),而不是收缩用的 tcontract;traced 和 eager 主线上,SVD、QR、solve 是扩展算子(family tenferro-linalg.linalg.v1),执行模块和求导规则都由 tenferro-linalg 提供。规则怎么注册、怎么分派见 semantic AD 与扩展,两种 provider 的选择和 native 对 Accelerate 的实测见 provider:native 与 BLAS。本页讲中间这一层和规则本身。

flowchart TB
  typed["a.svd(engine)"] --> wlb["with_linalg_backend"]
  traced["traced svd"] --> op["LinalgExtensionOp"]
  eager["eager svd"] --> op
  op -->|"runtime.run"| exe["execute_linalg"]
  wlb --> lb["LinalgBackend"]
  exe --> lb
  lb --> cpu["CpuEngine"]
  lb --> cuda["CudaEngine"]
  cpu -->|native| tl["tlinalg faer"]
  cpu -->|blas| tlb["tlinalg-blas"]
  cuda --> cus["cuSOLVER"]

tenferro-linalg 里有:backend.rs(后端 trait)、cpu/(CPU 适配:backend.rs 按 dtype 分派,linalg/faer_linalg.rs 和 linalg/lapack_linalg/ 是两种 provider 的宿主一侧,tlinalg.rs、tlinalg_blas.rs 是适配器)、gpu/(CUDA,只在 cuda feature 下编译)、extension.rs(扩展算子和执行模块)、traced.rs(traced 入口和组合运算)、ad/(求导规则)。数值内核和批循环都不在这里,在单独发布的 tlinalg 0.8 和 tlinalg-blas 0.5 里(仓库 tensor4all/tlinalg-rs)。实验在 exp-tenferro8/ops2-exp(la_misc、svd_timing、linalg_fd、svd_grad、dyn_trunc_grad、trunc_svd),CPU,全部 release,M2 Max、macOS 27、rustc 1.99.0。

有哪些运算

用户看到的是三组扩展 trait,方法名基本一致:

入口在哪执行
TypedTensorLinalgExt<T>、TensorLinalgExt、TensorReadLinalgExttensor_ext.rs立即执行,传 &mut dyn Engine
TracedTensorLinalgExttraced.rs只记图,返回 Result<TracedTensor>
EagerSessionLinalgExteager_ext.rs(autodiff 下)在借来的 eager 会话里执行并记录

traced 那一组有 25 个方法:svd、svd_with_options、svd_full、qr、householder_qr、qr_with_options、rank_revealing_qr、eigh、eigh_with_options、cholesky、lu、full_piv_lu、eig、solve、lstsq、full_piv_lu_solve、triangular_solve、slogdet、det、inv、eigvalsh、eigvals、pinv、pinv_with_rtol、norm。矩阵都按列主序约定取形状:前两维是矩阵,其余是批维。

一部分在图里就是一个扩展算子,即 LinalgOp 的 28 个变体之一(有几个只在求导程序里出现,如 HouseholderQrAppendTangent);另一部分是组合运算,在 traced.rs 里用别的算子拼出来:

组合运算拼法
invsolve(a, I)
slogdet先校验方阵,再 LuFactor,从同一个分解出 SignDetFromLuFactor、LogAbsDetFromLuFactor
detsign * exp(logabsdet),注释说照搬 JAX,直接连乘 LU 对角线在动态范围大时会先溢出
lstsq薄 QR,再 triangular_solve(R, Qᴴ b)
pinv_with_rtolsvd,按 rtol * max(s) 做掩码求 1/s,再 V diag(1/s) Uᴴ

组合运算的导数来自组成它的各个算子的规则,不需要单独写。

上游 linalg.md 在 v0.8.0 上仍有三处和代码对不上:它列的组合运算 matrix_power、cond、tensorinv、tensorsolve、multi_dot、vecdot、vander 在仓库里一个也搜不到;开头的层次图说 tenferro-tensor 的「TensorBackend」负责线代执行,实际的后端 trait 是本 crate 的 LinalgBackend;「The former host lane plan and CpuBatchPolicy plumbing were removed by #2004」这句也过时了,CpuBatchPolicy 在 0.8.0 里又回来了(见 engine 与 threadteam 池)。

两种 QR 的变体

普通 qr 不选主元,返回薄的 (Q, R)。另有两个为自适应算法加的运算:

  • 秩揭示 QR(rank_revealing_qr,选项 RankRevealingQrOptions,设计见 rank-revealing-qr.md):列主元 QR,返回 Q、R、列置换(I64)和数值秩(I64 标量),秩是 |R[i,i]| > max(atol, rtol·|R[0,0]|) 的前缀长度。秩是张量而不是 usize,输出形状永远是满的薄形状,不随数据变。主元选择和秩都不连续,所有输出都不可微。
  • 增量 Householder QR(HouseholderQr<T>,设计见 incremental-householder-qr.md):状态是压缩的反射向量,append_columns 追加一块列时只把已有反射作用到新块上,不重新分解,也不显式构造 Q;T 可以是 Tensor、EagerTensor、TracedTensor,满秩区域上有 JVP 和 VJP。

full_piv_lu 现在在两种构建里都是同一个 faer 内核,对奇异矩阵也能分解,来历见 provider 页的冷知识。

怎么落到 faer

CpuEngine 自己就实现了 LinalgBackend(trait LinalgBackend: Engine),不再有会话这一层。直接调用的入口都经过 with_linalg_backend:它对传进来的 &mut dyn Engine 先探测是不是 CpuEngine,再在 cuda feature 下探测 CudaEngine,都不是就报 selected backend session does not expose a linalg execution capability。traced 和 eager 走 define_extension_runtime! 生成的执行模块,execute_linalg 按 LinalgOp 分派到同一个 trait 的方法上。两条路的终点都是 impl LinalgBackend for CpuEngine。

每个方法的形状都一样:with_linalg_pool_checked 借出 engine 的执行令牌(tcontract_exec::Exec,里面是 threadteam 的令牌和批处理策略)和缓冲池,BLAS 构建先拒绝强制的批处理策略;然后 svd_entered 这类函数按 dtype(F32、F64、C32、C64)分派到 linalg::faer::svd 或 linalg::blas::svd。native 和 blas 是互斥的 feature,所以每个函数体里只有一支被编译进来。

faer 那一侧的宿主代码在 faer_linalg.rs,它的模块注释把分工说得很清楚:「Every numerical kernel and the batch loop live in the extracted tlinalg crate」,留在 tenferro 的是形状校验、空矩阵、池化输出、描述符和张量构造。SVD 是 svd_impl:

    let mut u = pooled_output::<T>(buffers, "svd", "left singular vectors", &u_shape)?;
    let mut s = pooled_output::<T>(buffers, "svd", "singular values", &s_shape)?;
    let mut vt = pooled_output::<T>(buffers, "svd", "right singular vectors", &vt_shape)?;
    let mut workspace = pooled_workspace::<T>(
        buffers,
        tlinalg::svd::svd_workspace::<T>(m, n, full, batch, exec).map_err(provider(Op::Svd)),
    )?;
    tlinalg::svd::svd_into(
        Op::Svd,
        raw_view(op, view)?,
        full,
        u.as_uninit_slice_mut(),
        s.as_uninit_slice_mut(),
        vt.as_uninit_slice_mut(),
        workspace.bytes_mut(),
        exec,
    )
    .map_err(provider(Op::Svd))?;

三个输出和工作区都从缓冲池取。tlinalg 的每个运算族都有一对接口:*_workspace 用同一个 Exec 算出这次调用要多少字节(按它会分几条批 lane 算),*_into 把每条 lane 的 faer 临时空间从这块字节里切出来。PooledWorkspace 就是池里一块缓冲的字节视图,调用结束就回池;输出用 assume_init_recycled 交出,drop 时也回池(机制见逐元素、SIMD 与融合)。并行同样由令牌决定:批里的矩阵数不少于令牌的宽度时每个线程串行算几个,否则每个矩阵拿一份宽度交给 faer 的 Par::Team,串行 engine 一律 Par::Seq(tlinalg 的 crate 文档「Parallelism」一节)。

tlinalg 交回的是 Vᴴ 而不是 V,转置在 provider 里做(svd_item 的最后一个循环)。它还多做了一件 faer 自己不做的事。faer 的默认参数在较小一维达到 128 时改用分治的双对角 SVD;svd.rs 的注释说这条路「can return inaccurate factors without an error」,所以带向量的分解在这个尺寸以上先用默认参数算,再用一次矩阵乘检查 ‖A − U diag(S) Vᴴ‖ 是否在 residual_bound 以内,不合格再用 qr_params(永不分治)重算:

    let accurate = match scratch.check.as_mut() {
        None => false,
        Some(check) => {
            faer::linalg::svd::svd(
                // ... 默认参数
            )
            .is_ok()
                && reproduces(
                    mat,
                    MatRef::from_column_major_slice(&*u_region, m, u_cols),
                    scratch.s.as_ref(),
                    scratch.v.as_ref(),
                    check,
                    <T as ScalarEntity>::EPSILON,
                    par.reborrow(),
                )
        }
    };

只要奇异值的路径(svd_values)不检查,一律用 QR 算法。这个检查是 tlinalg 的提交 9d94a32「fix: check the divide-and-conquer SVD and repeat with QR on failure」加的,注释点名的是 faer 0.24.4,而 tenferro 0.8.0 链接的是 t4a-faer 0.27。它的代价在下面的计时里。

LAPACK。 blas 构建把同样的调用交给 tlinalg-blas 0.5:SVD 是 ?gesdd(provider-inject 时换成 ?gesvd),工作区经 TlinalgWorkspace 也从同一个池里取。它不接收并行令牌,tlinalg_blas.rs 的注释说批循环是串行的,「the vendor call inside it owns its own threading」。0.7.1 时这条路在本机离线编不出来,0.8.0 的 blas-accelerate 能编能跑,两种 provider 的 SVD、QR 实测在 provider 页。

每次调用分配几次

#2188(首次进入 v0.8.0)把 faer 这条路改成上面的借工作区接口,提交说明的结论是「the faer lane's steady state drops to 3 allocations per call at one thread」,上游测试 cpu_linalg_allocation.rs 把 48×48 的 svd 钉在 3 次。la_misc 第 4 段用一个计数的全局分配器,在串行 engine 上对 f64 方阵各调 10 次取平均(先预热 3 次):

n=   8:  svd 2.0 (0 <=64B)  qr 2.0 (0 <=64B)  eigh 2.0 (0 <=64B)  solve 0.0 (0 <=64B)  svdvals 0.0 (0 <=64B)  matmul 0.0 (0 <=64B)
n=  32:  svd 2.0 (0 <=64B)  qr 2.0 (0 <=64B)  eigh 2.0 (0 <=64B)  solve 0.0 (0 <=64B)  svdvals 0.0 (0 <=64B)  matmul 2.0 (0 <=64B)
n=  48:  svd 33.0 (0 <=64B)  qr 29.0 (0 <=64B)  eigh 18.0 (0 <=64B)  solve 3.0 (0 <=64B)  svdvals 0.0 (0 <=64B)  matmul 2.0 (0 <=64B)
n=  64:  svd 49.0 (0 <=64B)  qr 45.0 (0 <=64B)  eigh 26.0 (0 <=64B)  solve 5.0 (0 <=64B)  svdvals 0.0 (0 <=64B)  matmul 2.0 (0 <=64B)
n= 256:  svd 412.0 (149 <=64B)  qr 403.0 (120 <=64B)  eigh 197.0 (64 <=64B)  solve 60.0 (11 <=64B)  svdvals 95.0 (16 <=64B)  matmul 2.0 (1 <=64B)

32 以下和上游的数字一致,48 起每次 SVD 有几十次分配,用上游测试里的同一个矩阵也是 33 次。alloc_trace 给每种大小的分配打一次调用栈,全部来自同一处:

size 1024:
   2: <dyn_stack::mem::MemBuffer>::new
   3: gemm_f64::gemm::f64::neon::gemm_basic
   4: faer::linalg::matmul::matmul_imp::<f64>
...

faer 分块算法里的矩阵乘在 aarch64 上走 gemm 0.19 这个 crate,它的 gemm_basic 每次调用都 MemBuffer::new 一块打包缓冲;t4a-faer 的 Cargo.toml 只在 x86_64 上引入另一个 GEMM 实现。所以「3 次」是 x86 上的数字(上游的测量机是 x86),在这台 M2 Max 上,矩阵大到 faer 开始分块之后,tenferro 自己的分配仍然是那几次,faer 内部每个 GEMM 还要各分配一两次。上游测试在 aarch64 上大概会失败,这里没有跑。

只用奇异值时换成更便宜的算子

traced 程序编译时会裁剪扩展算子没用到的输出(见 编译)。linalg 的 prune_outputs 只认三种情形:

            LinalgOp::Svd {
                derivative_eps,
                driver,
                ..
            } if live_outputs == [false, true, false] => {
                Some(Arc::new(Self::new(LinalgOp::SvdVals {
                    derivative_eps,
                    driver,
                })))
            }

Svd 只剩 S 活着时换成 SvdVals,Eigh、Eig 只剩特征值时同理。la_misc 第 3 段打印编译后程序里的算子:

S only:  ["SvdVals"]
U only:  ["Svd"]
Vt only: ["Svd"]
eigh values only: ["Transpose", "DotGeneral", "EighVals"]

只要 U 的程序照样算出 Vᴴ 再丢掉。还有一个副作用:SvdVals 走的是 faer 不算向量的那条路,和带向量的分解在最后一位上可以不同。下文 svd_grad 用的那个 4×3 矩阵,只取 S 时第二个奇异值是 1.9999999999999996,让 U 也活着时是 1.9999999999999993。

实测:比直接调 faer 慢多少

svd_timing 用同一个方阵(xorshift 生成的 [-0.5, 0.5) 均匀分布,f64)比较:tenferro 在 CpuEngine::serial() 上调 a.svd(..);faer 直接调 tenferro_cpu::faer::linalg::svd::svd(Par::Seq),分别用默认参数和只用 QR 算法的参数;再加一个 8 线程 engine 上的 tenferro。每项预热 1 次、计 9 次取最小值,跑了三轮,uptime 负载 6 左右(别的写手在跑程序),下表取三轮各自的最小值(ms):

ntenferro 串行faer 默认faer 只用 QRtenferro 8 线程
640.4470.4380.4370.429
1272.2882.2652.2462.369
1281.9091.7172.2471.579
25610.8519.93415.2987.674
51263.59557.563101.61530.332
n= 512: tenferro serial   63.595 ms  faer default   58.525  faer QR-only  101.615  tenferro team8   30.332  max|ds| vs default 0.0e0 vs QR-only 1.1e-12
  • 127 以下 tenferro 和直接调 faer 只差 1%–2%,没有检查,也不分治。
  • 128 起差 9%–11%,多出来的正是那次残差检查:512 时差约 6 ms,一次 512³ 的串行 f64 矩阵乘按 provider 页 的每核 48 GFLOP/s 算是 5.6 ms。这一次检查换来的是分治失败时能退回 QR 算法,而 QR 算法本身在 512 时要慢 1.7 倍。
  • 奇异值和 faer 默认参数逐位相同(max|ds| 0.0e0),和只用 QR 的结果差 1e-12 量级:tenferro 没有在数值上做任何加工,检查通过就用分治的结果。
  • 8 线程在 512 时是 30–35 ms,比 provider 页的 25.0 ms 慢,那一轮的负载是 2.2 到 2.7,这里是 6;串行的 63.6–66.2 ms 和那边的 64.9 ms 一致。

det 只接受方阵

la_misc 第 1、2 段对一个 2×3 矩阵在四个面上求行列式,再看奇异矩阵和空矩阵:

typed det: Err(det_read: incompatible shapes: lhs=[2], rhs=[3])
typed slogdet: Err(slogdet_read: incompatible shapes: lhs=[2], rhs=[3])
traced det (concrete shape): Err(slogdet: incompatible shapes: lhs=[2], rhs=[3])
traced det, symbolic input bound to [2, 3] at compile: Err(extension family "tenferro-linalg.linalg.v1" shape constraint at instruction Some(1) failed: 2 (2) Equal 3 (3))
eager det: Err(slogdet: incompatible shapes: lhs=[2], rhs=[3])
typed lu of the same 2x3: ok, U shape [[2, 2], [2, 2], [2, 3]]
det singular: Ok(Ok([1.3322676295501906e-15]))
slogdet singular: Ok((Ok([1.0]), Ok([-34.2518939198891])))
det 0x0: Ok(Ok([1.0]))
solve 0x2: Err(solve: incompatible shapes: lhs=[0], rhs=[2])

形状已知时在建图那一刻由 validate_square_matrix 拒绝;形状是符号时 slogdet 记下一条「行数等于列数」的符号约束,绑定 [2, 3] 编译时失败。lu 本身仍然接受长方形矩阵。奇异矩阵不报错:部分主元 LU 不把奇异当错误(tlinalg 的约定表),行列式给出舍入误差量级的 1.3e-15,slogdet 的对数是 −34.3。0×0 的行列式是 1,空的长方形 solve 被拒绝。

📝NOTE
冷知识:长方形矩阵的行列式曾经有值

PR #2329(提交 feaf32b3,首次进入 v0.8.0,发布当天合并)之前,det 和 slogdet 直接调能处理长方形矩阵的 LU,再把前 min(rows, cols) 个主元乘起来。提交说明的原话是「a rectangular input returned a plausible but wrong scalar on every surface」。修法是四个面共用一个先校验秩和方阵的辅助函数,traced 的行列式元数据记下符号方阵约束(上面那条编译期错误),LU 保持长方形。顺手审计其余只接受方阵的运算时还发现,solve、inv、full_piv_lu_solve 的空矩阵快速路径会对空的长方形输入返回一个值,也一并改成先校验(上面最后一行)。

AD 规则:先线性化,再转置

linalg 只注册了 linearize 和 linear_transpose 两种角色,没有 primal_vjp,两者都由 LinalgAdRule 实现,经 semantic.rs 的适配层翻译成语义程序(适配层见 semantic AD 与扩展)。所以反向一律是「先线性化,再转置」。

linearize 按 LinalgOp 分派到 ad/rules/ 里的函数,每个函数用核心算子(dot_general、逐元素、extract_diag、embed_diag……)把切向量的公式写进程序。linear_transpose 的 match 却大半是空的:

            // The fused solve is not linear in `a`; its linearization emits
            // `LuSolvePrepared`, whose transpose rule above handles VJP.
            LinalgOp::Cholesky
            | LinalgOp::Lu
            // ...
            | LinalgOp::Svd { .. }
            | LinalgOp::SvdFull
            | LinalgOp::SvdVals { .. }
            | LinalgOp::Qr { .. }
            // ...
            | LinalgOp::HouseholderQrSplitTangent { .. } => Ok(vec![None; op.input_count()]),

真正写了转置的只有对某个输入线性的几个:TriangularSolve、LuSolvePrepared、FullPivLuSolve、Solve、HouseholderQrAppendTangent。这一块从 0.7.1 到 0.8.0 没有变:ad/rules/mod.rs 在两个 tag 之间没有改动,ad.rs 只改了测试。

📝NOTE
没有蠢问题

问: Svd 的转置规则返回全 None,那 SVD 的反向不就是零吗?

答: 转置作用在线性化之后的程序上,那里面根本没有 Svd。linearize_svd 发出的切向量程序只由 dot_general、逐元素乘除、对角线抽取这些核心算子组成,U、S、Vᴴ 是作为常数(残差)进去的;反向扫描转置的是这些核心算子,用的是核心规则。只有当线性化的结果里还留着一个 linalg 算子时,才需要 linalg 自己的转置:solve 的切向量是「用已分解的 A 再解一次」,这一步就是 LuSolvePrepared,它对右端项线性,所以要写转置。返回 None 的分支在正常的反向中不会被调用。

规则不支持的运算会报错,而不是给零梯度。linalg_fd 对秩揭示 QR 的 R 和 svd_full 的奇异值求导:

rrqr grad: Err(vjp (GraphBuild): runtime state failure: semantic extension AD failed: semantic extension family "tenferro-linalg.linalg.v1" Linearize rule failed: unsupported AD rule Jvp for tenferro-linalg.rank_revealing_qr)
svd_full sum(s) grad: Err(unsupported vjp AD rule for tenferro-linalg.linalg.v1)

哪些运算、哪些输出可微,可以在运行期查 linalg_ad_support,秩揭示 QR 那一条带着理由:caveats: ["Pivot selection and numerical rank are discontinuous; all outputs are initially unsupported for AD."]。

有限差分核对

linalg_fd 对五种运算各构造一个标量损失,用带 linalg 规则的 AdContext 求梯度,再和中心差分(步长 1e-6)比较。Cholesky 的输入是 X Xᵀ + 3I、eigh 的输入是 X + Xᵀ,对 X 求导,这样扰动不破坏前提;eigh 的损失包括特征值和第三个特征向量的投影矩阵 v₃ v₃ᵀ(与符号无关):

solve  loss=-0.759831 max|grad|=0.867 max|grad-fd|=7.852e-11
chol   loss=+0.279193 max|grad|=0.517 max|grad-fd|=2.762e-10
qr     loss=-0.979951 max|grad|=1.157 max|grad-fd|=3.210e-10
eigh   loss=-2.630406 max|grad|=1.363 max|grad-fd|=1.220e-9
det    loss=+2.131000 max|grad|=1.660 max|grad-fd|=4.504e-10

同一个损失(U、S、Vᴴ 各自加权求和,4×3)在 traced 下用 AdContext::grad、在 eager 下用 EagerRuntime::with_ad_context 的会话调 backward,结果逐位相同(bitwise equal: true):两边用的是同一条规则、同一个变换。

solve 只分解一次

traced 的 solve 只发一个算子:

pub fn solve(a: &TracedTensor, b: &TracedTensor) -> Result<TracedTensor> {
    // One fused op factors and solves in a single backend call and saves
    // (x, packed LU, pivots), so AD reuses the factors for the adjoint solve.
    let mut outputs = apply(
        Arc::new(LinalgExtensionOp::new(LinalgOp::LuFactorSolve)),
        &[a, b],
    )?
    .into_iter();

LuFactorSolve 有三个输出 (x, packed_lu, pivots),用户只拿到 x,另外两个是给反向留的:线性化发出 LuSolvePrepared,直接用存下的分解解伴随方程,反向模式只分解 A 一次。prune_outputs 里有一段注释说明它为什么不能被裁成只有一个输出的 Solve:traced AD 在求导前先按输出裁剪源程序,一旦裁掉,分解就没了,伴随求解得重新分解 A。

📝NOTE
冷知识:solve 从两个算子合成了一个

PR #1880(「perf(linalg): tight batched CPU LAPACK loops and fused solve lowering」,首次发布于 v0.7.0)之前,traced solve 先发一个 LuFactor,再拿它的 packed_lu 和 pivots 发一个 LuSolvePrepared,一次求解是两个扩展算子、两次后端调用。提交说明:「Traced solve and tracked eager solve emit one LuFactorSolve op with outputs (x, packed LU, pivots). The CPU kernel fuses factor and solve; the linearize and transpose rules reuse the saved factors, so reverse mode factors A once. LuFactorSolve is never pruned to Solve because traced AD prunes the source program before differentiating it.」0.8.0 里这个融合内核搬进了 tlinalg(packed_lu::factor_solve_into),tenferro 这边的 factor_solve_batch 只负责给它借工作区。

SVD 的反向规则

linearize_svd 写的是薄 SVD A = U diag(s) Vᴴ(A 是 m×n,k = min(m, n))的切向量。记 dP = Uᴴ dA V(k×k),它把公式拆成几块:

  • 奇异值:ds = diag(dP),只有 S 被用到时到这里就返回。
  • 两两耦合的系数矩阵 F,F[i,j] = (s_j² − s_i²) / ((s_j² − s_i²)² + ε²)。
  • dU = U (F ∘ (dP S + (dP S)ᴴ) + ½ diag(1/s) ∘ (dP − dPᴴ)),最后一项是 JAX 的复数规范修正,实数时为零;m > n(或编译时形状未知)时再加 (I − U Uᴴ) dA V S⁻¹。
  • dV = V (F ∘ (S dP + (S dP)ᴴ)),n > m(或形状未知)时再加 (I − V Vᴴ) dAᴴ U S⁻¹。

F 那一段:

    let s_sum = fixed_add(builder, ValueRef::Local(s_dim), ValueRef::Local(s_dim_t));
    let s_diff = fixed_sub(builder, ValueRef::Local(s_dim), ValueRef::Local(s_dim_t));
    let s_gap = fixed_mul(builder, ValueRef::Local(s_sum), ValueRef::Local(s_diff));
    let s_gap_sq = fixed_mul(builder, ValueRef::Local(s_gap), ValueRef::Local(s_gap));
    let eps_sq = fixed_scale_with_dtype(
        // ... ones * eps * eps
    )?;
    let safe_gap = fixed_add(builder, ValueRef::Local(s_gap_sq), ValueRef::Local(eps_sq));
    let f = fixed_div(builder, ValueRef::Local(s_gap), ValueRef::Local(safe_gap));

教科书的写法是 1 / (s_j² − s_i²),两个奇异值相等时除以零。这里乘上 gap / gap 再在分母加 ε²:间隙远大于 ε 时和原式一样,间隙为零时 F 恰好是 0,中间是一条洛伦兹形的过渡,峰值 1/(2ε)。diag(1/s) 也写成 s / (s² + ε²)。ε 是 SvdOptions::derivative_eps,默认 DEFAULT_DECOMPOSITION_DERIVATIVE_EPS = 1e-12,是算子载荷的一部分,前向执行时只做合法性检查。eigh 的特征向量规则用的是同一种展宽。

只有两处没有展宽:m > n 和 n > m 时那一项投影修正,linear_div_fixed(builder, proj, s_broadcast) 直接除以原始的 s。

简并和零奇异值的实测

svd_grad 用固定的 Householder 反射拼出 A = Q[:, :3] diag(s) Rᵀ(4×3,m > n),让奇异向量不是单位向量,再对几种损失求导、和中心差分比较。损失有三种:Σs;P_k = U_k U_kᴴ 加权求和(前 k 个左奇异向量张成的投影,s_k > s_{k+1} 时与规范无关);U[0,0](依赖符号,是个病态的函数)。

输入的 s损失max |grad − fd|说明
(3, 2, 1)四种都试3.7e-10 – 8.1e-10普通情形,全部对
(2, 2, 1)Σs7.3e-10核范数在满秩处可微
(2, 2, 1)P_21.1e-6投影到简并块整体,梯度有定义,也对
(2, 1, 0),经反射P_25.8e-10faer 算出的 s₃ 是 3.4e-16,不是 0
(2, 1, 0),对角矩阵P_2全是 NaN差分是有限的(-0.6、0.7、1.4 等)
3×3 对角 (2, 1, 0)P_26.1e-11方阵没有投影修正项

对角输入的那一行:

diag_rank2 Proj(2) eps=None: s = [2.000000, 1.000000, 0.000000]
  grad = [NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN, NaN]
  fd   = [0.000000, 0.000000, -0.600000, 0.700000, 0.000000, 0.000000, 1.400000, -0.400000, 0.000000, 0.000000, 0.000000, 0.000000]

这和 0.7.1 的结果一样,是一个仍在的上游问题。对角输入让 faer 给出精确的 s₃ = 0,m > n 的修正项 proj / s 得到 0/0 或 x/0;即使 P_2 用掩码把第三列乘了 0,0 · inf 仍是 NaN,再经矩阵乘扩散到梯度的每个元素。同样的奇异值换成方阵就没有这一项,梯度和差分一致。真实的秩亏损矩阵会给出 1e-16 量级的 s₃,这时梯度有限;但若损失真的依赖第三个奇异向量,修正项会被放大到 1e15 量级(同一输入的 P_3 = U Uᴴ,梯度 4e15、差分 1e-3,这个函数在秩亏损处本身就不可微)。

❗IMPORTANT
动脑筋:奇异值是 (2, 2, 1),损失是只依赖第一个左奇异向量的 P_1 加权和,默认 ε = 1e-12。梯度会是 NaN、报错,还是别的什么?

提示:faer 算出的两个 2 不会逐位相等,差多少?代进 F 的式子。

换一个 ε

derivative_eps 可以按调用改:x.svd_with_options(SvdOptions::default().derivative_eps(1e-6))。同一个 (2, 2, 1) 输入,ε 取 1e-6 和 1e-3 时,P_2 的梯度与差分的误差是 3.4e-10 和 5.5e-8;(3, 2, 1) 的 P_1 是 6.3e-10 和 5.8e-9。间隙为 1 时,ε 大到 1e-3 也几乎不影响正确的导数;ε 改变的只是间隙小于 ε 的那些耦合,把它们压向 0。默认 ε 下 (2, 2, 1) 的 P_2 误差是 1.1e-6,比别的行大,很可能是块内那个 1e9 量级的系数(见动脑筋答案)在对称的两项之间没有完全抵消。

截断 SVD 与动态形状

上游例子 dynamic_shape_truncated_svd 演示「编译一次,按运行期的数据决定保留几个奇异值」:keep_count = Σ (s > 0.5) 是图里的一个标量,s.dynamic_truncate(&keep_count, 0) 的形状元数据是上界(UpperBound,见 TracedTensor),执行时才变成 2 或 3。它在 base 和 ops2-exp 里 release 下都以 0 退出。

细看会发现它只截了 s。重建矩阵那一段把 s 乘一个 0/1 掩码,U、Vᴴ 保持 4×4,乘出来的还是静态形状。实验 trunc_svd 改成真正的截断:U、s、Vᴴ 都按 k 截(U 截第 1 维),再用 broadcast_in_dim_sym 和 dot_general 乘回 4×4,损失是加权和。前向没有问题,梯度程序能建、能编译,执行时失败:

diag [4.0, 3.0, 0.1, 0.01]: s_k = [3.9999999999999996, 3.0000000000000004] shape [2]
  loss = [-2.3262222222222224]
  grad run: Err(TypedTensorView::reshape_view: reshape element-count mismatch: from 8 to 16)

换回例子里的掩码写法,同一个损失的梯度与差分差 2.2e-9(k = 2)和 2.6e-9(k = 3),程序只编译一次。

失败和 SVD 无关。dyn_trunc_grad 把它缩到最小:一个常数张量 x,k = 2,x_k = x.dynamic_truncate(&k, axis),对三种损失用 TracedTensorAdExt::grad,不涉及任何扩展算子:

x [5] axis 0 sum(x_k * x_k): grad [5] [2.0, 4.0, 0.0, 0.0, 0.0]
x [4, 3] axis 0 sum(x_k): grad [4, 3] [1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0]
x [4, 3] axis 0 sum(x_k * x_k): Err(TypedTensorView::reshape_view: reshape element-count mismatch: from 6 to 12)
x [4, 3] axis 0 sum(exp(x_k)): Err(TypedTensorView::reshape_view: reshape element-count mismatch: from 6 to 12)
x [4, 3] axis 1 sum(x_k * x_k): Err(TypedTensorView::reshape_view: reshape element-count mismatch: from 8 to 12)
x [1, 5] axis 1 sum(x_k * x_k): Err(TypedTensorView::reshape_view: reshape element-count mismatch: from 2 to 5)

一维输入全部正确;二维输入只要截断后的值被一个「反向要读它」的运算用到(乘、exp),就在执行时失败,sum(x_k) 这种不读它的则没问题,连 [1, 5] 这种只多一个长度为 1 的轴的也失败。ReduceSum 的 VJP 专门调了 truncate_ad_value_to_dynamic_axes 处理动态轴,乘法这一类没有;上游的 dynamic_truncate_vjp_correct 测试仍然只覆盖一维的 x * x。和 0.7.1 比,失败的位置和条件没变,报错从 reshape: incompatible shapes: lhs=[2, 3], rhs=[4, 3] 换成了按元素数报的 reshape_view。要对截断 SVD 求导,仍然只能用掩码写法。

CUDA(只读源码)

cuda feature 下 CudaEngine 也实现 LinalgBackend,代码在 gpu/,经 cuSOLVER。0.8.0 release notes 的已知限制里写着 CUDA 不支持 eig、full_piv_lu、full_piv_lu_solve(返回 Unsupported);CUDA 的 solve 不再为了检查奇异性把 LU 状态读回主机(#2204),精确奇异的输入得到 provider 定义的值,lu_factor 和 lu 仍报 Singular。本机没有 CUDA,这些只来自源码和 release notes,engine 本身见 CUDA。

动脑筋答案

svd_grad 的输出:

repeated Proj(1) eps=None: s = [2.000000, 2.000000, 1.000000]
  grad = [-2124177597.858788, 125079545.236988, -949734675.103228, -2849204025.357783, 1590289147.672616, 775068093.833715, 16061183.927504, 48183551.798115, -2223750359.331037, -1896047900.547346, 627338176.718546, 1882014530.152900]
  fd   = [-512661.915492, -345673.088298, -125143.091501, -125143.177894, 305425.333679, 505853.072401, -357985.739621, -357985.980936, 16869.578293, -339276.884417, 509834.876191, 509835.309739]
  max |grad - fd| = 2.849e9
repeated (with vectors): s = ["2e0", "1.9999999999999993e0", "1.0000000000000007e0"]
  • 既不是 NaN 也不报错,是一个 1e9 量级的有限数。带向量的分解给出的两个奇异值是 2e0 和 1.9999999999999993e0,s_j² − s_i² 约 2.8e-15;它的平方约 8e-30,远小于 ε² = 1e-24,于是 F ≈ gap / ε² ≈ 2.8e9,正是梯度的量级。
  • 展宽在 gap = ε 时才到峰值 1/(2ε) = 5e11,间隙比 ε 小时 F 正比于间隙,间隙精确为 0 时才是 0。对角输入正好是这种情况:同一个损失换成 diag(2, 2, 1),梯度是有限的小数,块内的方向给 0。
  • 差分也是 1e5 量级、而且没有意义:简并块里「第一个」奇异向量是哪个由舍入决定,P_1 在这一点不可微。规则给出的大数和差分都不是导数。
  • 把 ε 调到 1e-6 后,同一个输入的梯度变成 0.1 量级(-0.090754, 0.198083, ...),这是把块内耦合压到 0 之后的结果,仍然不是导数,只是不再爆炸。

小结

  • CpuEngine 直接实现 LinalgBackend;直接调用经 with_linalg_backend,traced / eager 经扩展执行模块的 execute_linalg,都落到按 dtype 分派的 *_entered,再进 tlinalg(faer)或 tlinalg-blas(LAPACK)。上游 linalg.md 的组合运算清单和层次图与代码不符。
  • tlinalg 的 *_workspace / *_into 让工作区和输出都从 engine 的缓冲池借,32×32 以下每次 SVD 只分配 2 次;更大的矩阵在 aarch64 上每个 faer 内部 GEMM 还会分配打包缓冲,48×48 是 33 次。
  • 较小一维不小于 128 时,tlinalg 对 faer 的分治 SVD 做一次残差检查,不合格就用 QR 算法重算;这让 tenferro 比直接调 faer 慢约 10%,小于 128 时只差 1%–2%,奇异值逐位相同。
  • det / slogdet 在四个面上都拒绝非方阵(#2329,v0.8.0),符号形状在编译时失败;奇异矩阵给出舍入量级的行列式,不报错。
  • AD 只有 linearize 和 linear_transpose;solve、Cholesky、QR、eigh、det、SVD 与有限差分一致,eager 与 traced 逐位相同;秩揭示 QR、svd_full 报错而不是给零。SVD 规则用 gap / (gap² + ε²) 处理简并,近简并得 1e9 量级的数;m ≠ n 的投影修正直接除以 s,精确为 0 的奇异值让整块梯度变成 NaN,和 0.7.1 一样。
  • 二维以上 dynamic_truncate 之后的乘法类 VJP 仍然在执行时失败,报错换成了 reshape_view: reshape element-count mismatch;截断 SVD 求导用掩码写法。

💬 关联讨论

发起讨论
加载中...