什么时候用 tenferro

主线位置：TRG 主线的「分段求导得内能」，加上变换主线的「编译、准备 → 执行」。本页比较同样的任务两边多快、交付出去有多大、启动要多久，再列出该留在 JAX 的情况。

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

flowchart TB
  task["同一个任务"] --> jx["JAX jit"]
  task --> tf["tenferro"]
  jx --> jfast["稳态快约 2 倍"]
  jx --> jbig["环境约 409 MiB"]
  tf --> tsmall["二进制 21.3 MiB"]
  tf --> tstart["启动 0.5 ms"]
  tf --> tstage["长图要手工分段"]
  jx --> jhas["vmap scan 随机数"]

这一页只写在本机测得到的东西。平台 M2 Max（8 个性能核、4 个能效核）、macOS 27、rustc 1.99.0，tenferro 全部 release 构建。计时时机器上还有别的任务，每组数字旁边写了当时 uptime 的 1 分钟负载。JAX 每个函数先调用一次（含追踪和编译），再计稳态，每次都 block_until_ready()。

TRG：20 步，求值加梯度

任务是这一卷的主线：二维 Ising 模型，χ = 16，20 步，五个 β，求 lnZ/N 和 d(lnZ/N)/dβ。JAX 一侧把整个 20 步的函数交给 jit(value_and_grad)：

ks = schedule(16, 20, "rank")      # [2, 4, 4, 16, 16, ..., 16]
f = make_trg(ks, SVD["lorentz"])   # beta -> lnZ/N，Python 循环展开 20 步
fv = jax.jit(f)
fvg = jax.jit(jax.value_and_grad(f))

make_trg 的全文和 SVD 求导规则见动手：Ising TRG。tenferro 不能这样做：对整张 20 步的图求导，变换本身的时间随步数指数增长（下面「求导变换本身的代价」一节）。可行的写法是按每一步的形状 (d, k) 各建一个小 traced 程序，五种形状，20 步和所有 β 共用；反向在步与步之间手工串起来：

for &k in ks {
  if !steps.contains_key(&(d, k)) {
    let t = TracedTensor::input_concrete_shape(DType::F64, &[d, d, d, d])?;
    let ct = TracedTensor::input_concrete_shape(DType::F64, &[k, k, k, k])?;
    let c = TracedTensor::input_concrete_shape(DType::F64, &[])?;
    let (tn, ln) = trg_step(&t, d, k, None)?;
    let l = (&(&ct * &tn)?.reduce_sum(None)? + &(&c * &ln)?)?;
    let g = ad.grad(&l, &t)?;
    steps.insert((d, k), StepProgs {
      fwd: mk(rt, eng, &tn, &[(&t, vec![d, d, d, d])])?,
      lognorm: mk(rt, eng, &ln, &[(&t, vec![d, d, d, d])])?,
      bwd: mk(rt, eng, &g, &[(&t, vec![d, d, d, d]), (&ct, vec![k, k, k, k]), (&c, vec![])])?,
    });
  }
  d = k;
}
// ...
for i in (0..n).rev() {
  let (d, k) = (dims[i], dims[i + 1]);
  let c = scalar(2f64.powi((n - i) as i32) * scale)?;
  ct = run(rt, eng, &s.steps[&(d, k)].bwd, &[&ts[i], &ct, &c])?;
}

mk 是编译加 prepare，run 是 runtime.run 加读回。这相当于 lax.scan 加逐步的 VJP，写法和取舍见分段。

两边的结果对得上：lnZ/N 最大差 5.6e-16，梯度最大差 2.4e-12，和 tenferro 自己换线程数的差同一量级（怎么读这一卷）。每个 β 的稳态时间，β 之间的范围（JAX 每个 β 取 10 次最小，tenferro 取 5 次最小）：

几点说明：

• JAX 不是单线程。 同一段稳态循环里进程的 CPU 时间除以墙钟时间，求值是 1.22–1.26，求值加梯度是 2.10–2.24。设 XLA_FLAGS="--xla_cpu_multi_thread_eigen=false intra_op_parallelism_threads=1" 再跑，时间和这个比例都没变（160.0–192.1 / 223.8–252.5 ms，比例 1.22 / 2.09–2.22），本页没有找到让它只用一个核的开关。所以拿它和 tenferro 的单线程、8 线程两行都比一比：求值和 8 线程相当，求值加梯度快约 2 倍。
• tenferro 的求值加梯度是求值的 2.5–2.7 倍，JAX 是 1.3–1.4 倍。 分段写法里每一步的反向程序 bwd 以这一步的输入 T 为参数，自己重算这一步的前向，两次 SVD 也重算；相当于每一步都做了一次 checkpoint。JAX 对整张图求导，前向的中间结果直接留给反向。
• 准备的代价正好反过来。 JAX 对 20 步的图追踪、求导、编译，第一次调用 1.2 秒；tenferro 只编五个小程序，18.5 ms，之后所有 β 复用（β 是占位符）。
• PyTorch 的梯度是 NaN。 torch.linalg.svd 的反向在这条链上给出 NaN；JAX 内置的 jnp.linalg.svd 在这条链上同样给出 NaN，上表的 JAX 一行是用 custom_jvp 补上了和 tenferro 默认一样的展宽规则。tenferro 不用补（简并奇异值）。

小函数的稳态

变换主线的 f(x) = Σ sin(x)·x²，n = 1000，每个程序跑 200 次取最小值，负载约 3.7：

eager 方式（每个运算立即执行，再反向）：tenferro with_eager_session 加 backward 73.5 µs，PyTorch 1 线程 29.4 µs，不加 jit 的 jax.grad 970.5 µs。

1000 个元素太小，8 线程比单线程慢：分活和等待的开销比计算本身大（12 个线程比 8 个慢）。小函数上 JAX jit 快 1.1 到 2.2 倍；eager 上 PyTorch 快约 2.5 倍。怎么计时、为什么要分开第一次调用，见计时。

求导变换本身的代价

上面 tenferro 要分段，是因为求导变换在长图上随深度指数增长。直接对 TRG 的前 n 步整体建图，单独计 ad.grad 这一步（只是变换，不含编译和执行），负载 3.7：

tenferro 0.8.0 traced chi=16 steps=2 schedule=rank width=1 ks=[2, 4]
ad.grad                         8.43 ms
tenferro 0.8.0 traced chi=16 steps=3 schedule=rank width=1 ks=[2, 4, 4]
ad.grad                        89.73 ms
tenferro 0.8.0 traced chi=16 steps=4 schedule=rank width=1 ks=[2, 4, 4, 16]
ad.grad                      2077.58 ms

每多一步，约 10 到 23 倍。把它缩到最小：3×3 矩阵上反复做 (u, s, vt) = svd(y); y = u·s + vt，三个输出都被用到，n 层：

svd3 n= 8 value compile     0.14 ms (57 ops)  ad.grad       2.93 ms  grad compile     1.40 ms (460 ops)
svd3 n=10 value compile     0.18 ms (71 ops)  ad.grad      12.45 ms  grad compile     1.72 ms (574 ops)
svd3 n=11 value compile     0.25 ms (78 ops)  ad.grad      32.95 ms  grad compile     1.79 ms (631 ops)
svd3 n=12 value compile     0.18 ms (85 ops)  ad.grad      93.31 ms  grad compile     1.72 ms (688 ops)
svd3 n=13 value compile     0.21 ms (92 ops)  ad.grad     275.46 ms  grad compile     2.12 ms (745 ops)

导数程序的大小（ops）随层数线性增长，ad.grad 的时间每层乘 2.6 到 3。实验只跑到 13 层，再往上每层要多等几倍。JAX 对整张 20 步的图做同样的事，加上 XLA 编译，第一次调用一共 1.2 秒。这个代价从哪来、分段之外还能怎么绕，见求导变换的代价。

交付出去有多大

JAX 一侧：Python 解释器，加 jax、jaxlib 和它们声明的依赖（ml_dtypes、NumPy、SciPy、opt_einsum）。按 du -sk -I __pycache__ 量，1 MiB = 1024 KiB：

这是 macOS 上只带 CPU 的 wheel。不算 PyTorch（它在同一个虚拟环境里，另占约 565 MiB）。

tenferro 一侧是一个可执行文件。上面那个 grad 程序只依赖 tenferro-runtime、tenferro-cpu、tenferro-ad；TRG 的 trg_staged 依赖全套（加 linalg、einsum、fft）：

strip 之后的 grad 程序照样输出同样的梯度。409.4 MiB 对 21.3 MiB，差约 19 倍；不像 candle 那样能小到 1 MiB 以下，tenferro 的二进制大在 strided-rs、threadteam 和 GEMM 微内核按 dtype、布局展开出来的代码（tenferro 卷的 tenferro 解决什么问题量过只做一次 matmul 的 30.3 MB）。

启动到第一个梯度

同一个 grad 任务做成一次性的进程：启动、建数据、第一次求梯度、打印、退出。JAX 一侧进程里第一行就开始计时：

import time
T0 = time.perf_counter()
import sys

import jax
T_IMPORT = time.perf_counter()
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp

N = 1000
x = jnp.array([-2 + 4 * i / (N - 1) for i in range(N)])
g = jax.block_until_ready(jax.jit(jax.grad(lambda x: jnp.sum(jnp.sin(x) * x**2)))(x))
T_GRAD = time.perf_counter()
# ... 打印 g，把两段时间写到 stderr

tenferro 一侧就是怎么读这一卷里那个程序，在 main 第一行取 Instant::now()。各跑 10 次，负载 2.5 到 2.7：

JAX 一行的整个进程时间来自 /usr/bin/time -l（分辨率 10 ms）。tenferro 在它下面显示 0.00 s，所以改用一个 Python 的 subprocess 计时器量；同一个计时器跑 /usr/bin/true 是 1.8 ms，这是它自己的底。（这个计时器量 Python 子进程会多出约 50 ms，python -c pass 也一样，原因没有追查，所以 JAX 一行不用它。）

JAX 的 0.28 秒里，import jax 占了大半，jit 第一次编译这个小函数约 35 ms。tenferro 在进程内从建图、求导、编译、准备到执行完一共不到 1 ms。按请求起进程、做命令行工具的场合，差别几乎全在这里；在一个长期运行的进程里反复求值，这笔一次性的账就摊薄了，比的是上面的稳态。

GPU

tenferro 0.8.0 唯一的 GPU 后端是 CUDA：tenferro-gpu 的 feature 只有 cuda（和一个空的 rocm），见 tenferro 卷的 CUDA。另有一个 tenferro-xla，把编译好的图翻成 StableHLO，交给外部的 PJRT 插件执行；本机没有插件，tenferro 卷也只读了源码（XLA）。所以在 Mac 上，tenferro 只有 CPU。

JAX 一侧，本卷的虚拟环境里 jax.devices() 只有 CpuDevice(id=0)，没有装任何 GPU 插件。这一卷两边都只在 CPU 上比。

JAX 有、tenferro 0.8.0 没有的

逐条在 tag v0.8.0 的源码上核对过：

• 没有 vmap。 crates/ 下搜不到 vmap，traced 张量上的变换只有 grad、vjp、jvp 和 checkpoint（TracedTensorAdExt）。批量要显式写进形状，比如 dot_general 的 batch 维。
• 没有 lax.scan、fori_loop、while_loop、cond。 核心算子表 StdTensorOp 里是算术、超越函数、DotGeneral、Reshape、Gather、Scatter、DynamicSlice、Pad 这类张量算子，加上 Extension，没有带子图的控制流算子。Rust 的循环在建图时展开，逐元素的选择用 select / where_select（相当于 jnp.where）。TRG 那种长链要像上面那样手工分段（分段）。
• 没有 jax.random。 TracedTensor 的方法和核心算子表里都没有随机数；要随机输入，在 Rust 里生成好数据再建张量。
• 没有 pytree。 求导的对象是一个个 TracedTensor：y.grad(&x) 对一个输入；多个输入用 AdContext::vjp_many 一次求出（tenferro 卷的 traced AD）。参数的嵌套结构要自己摊平。
• 没有函数级的 custom_jvp / custom_vjp。 不能给一段 traced 代码直接挂一条求导规则；自定义导数要把它写成扩展算子，再给这个算子族一条 semantic AD 规则（扩展算子、semantic AD 与扩展）。反过来，JAX 里要手写 custom_jvp 才有的 SVD 展宽规则，tenferro 是默认的。
• 长图求导的代价。 上面量过：JAX 对 20 步的图一次求导，tenferro 到 4 步就要 2 秒。

tenferro 这一侧也有 JAX 不同的地方，各页细讲：同一个库里既有 traced 变换也有 PyTorch 式的 backward（eager）；复数求导的约定和 PyTorch 一样，和 JAX 差一个共轭（复数求导）；张量是列主序（布局与 reshape）。

什么时候用哪个

• 用 tenferro：求导要放进 Rust 程序里交付（命令行工具、嵌进 Rust 服务、按请求起进程），在意启动时间、内存和部署体积：本页的数字是体积差约 19 倍，进程内到第一个梯度差约 500 倍（整个进程约 100 倍），峰值 RSS 差约 28 倍。用到 SVD、eigh 这类分解的梯度又有简并时，tenferro 默认就给出有限的梯度。数据本来就是列主序（Fortran、Julia、LAPACK）时也顺手。
• 留在 JAX：要 vmap、lax.scan、随机数、pytree 形式的参数；想对一条长计算链整体求导而不想手工分段；看重稳态速度（本页的 TRG 上，JAX 求值加梯度比 tenferro 8 线程快约 2 倍，小函数上快 1.1 到 2.2 倍）；在 notebook 里交互试验；要 GPU。

动脑筋答案

first exec of the stripped copy:
grad[0] = 1.972602361114157, sum(grad) = 1818.745425877291154
trace+grad+compile 6.104 ms, +runtime+prepare 9.722 ms, +first run 16.388 ms

16.4 ms，是之后的 30 多倍。

• 慢的不只是执行，三段都慢：建图、求导、编译 6.1 ms（之后 0.23 ms 左右）。这说明代码本身第一次被用到时有额外开销，不是 tenferro 的哪个缓存还没热。
• 实验里重新构建之后第一次运行也是这样（15.7 到 17.5 ms），隔了几分钟再跑一个旧的二进制也慢过一次（ad_vec，第一次 run 6.9 ms，之后 0.1 ms 左右）。本页没有追查是 macOS 的哪一步；可以确定的是第二次运行同一个文件就回到 0.5 ms 左右。
• 所以量冷启动时要说清楚是哪一种「冷」：新装的文件、久未运行的文件，还是刚运行过的文件。上面表里的 10 次都是第二次以后。

小结

• 20 步 TRG（χ = 16）：JAX jit 求值 161–190 ms、求值加梯度 223–253 ms；tenferro 必须分段，8 线程 177–189 / 441–472 ms，单线程 333–355 / 913–948 ms。JAX 的进程用了 1.2 到 2.2 个核。tenferro 的准备只要 18.5 ms，JAX 第一次调用 1.2 秒。
• 小函数（n = 1000）：JAX jit 的 grad、jvp、hvp 10–13 µs，tenferro 单线程 14–24 µs；eager 上 PyTorch 29 µs，tenferro 74 µs。
• tenferro 的求导变换在长图上随深度指数增长：TRG 2、3、4 步 8 ms、90 ms、2.1 s；3×3 的 SVD 链每层 ×2.6 到 ×3。
• 交付：JAX 加 Python 约 409 MiB，tenferro 的 grad 程序 strip 后 21.3 MiB，只链接系统库。到第一个梯度：JAX 进程 0.36 s 起，import jax 占大半；tenferro 进程 3.2 ms 起，进程内不到 1 ms；峰值 RSS 178 MiB 对 6.4 MiB。新文件第一次运行要 16 ms。
• tenferro 0.8.0 的 GPU 只有 CUDA，Mac 上只有 CPU；Metal 路径随 PR #2291 删掉了。
• 没有 vmap、控制流算子、随机数、pytree、函数级 custom_jvp；有默认的 SVD 展宽规则、eager backward 和 PyTorch 式的复数约定。