怎么读这一卷

主线位置：三条主线的总览。本页用变换主线的「grad → 编译、准备 → 执行」做例子，讲一对对照程序怎么跑、两边的数按什么标准算对上，再把三条主线的每一步对到讲它的页。

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

flowchart TB
  py["grad_jax.py"] --> pj["grad_jax.npy"]
  rs["tenferro 程序"] --> tj["grad_tf.txt"]
  pj --> cmp["compare_vec.py"]
  tj --> cmp
  cmp -->|"不超过容差"| ok["OK 退出码 0"]
  cmp -->|超出| bad["MISMATCH 退出码 1"]

锁定版本

平台：M2 Max（8 个性能核、4 个能效核）、macOS 27、rustc 1.99.0。

JAX 一侧，每个脚本在 import jax.numpy 之前打开 64 位：

import jax

jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp

不开的话，JAX 默认用 float32，还会把 f64 的输入静默降成 f32；tenferro 的 dtype 来自 Rust 的字面量，vec![1.0_f64] 就是 F64。对照 f64 的结果，两边必须都是 f64。默认值本身的差别见 dtype。

tenferro 一侧，每个实验 crate 的依赖是同一组：

[dependencies]
tenferro-runtime = "=0.8.0"
tenferro-cpu = "=0.8.0"
tenferro-tensor = "=0.8.0"
tenferro-ad = "=0.8.0"
tenferro-linalg = { version = "=0.8.0", features = ["autodiff"] }
tenferro-einsum = { version = "=0.8.0", features = ["autodiff"] }
tenferro-fft = { version = "=0.8.0", features = ["autodiff"] }

线性代数、einsum、FFT 的求导规则在各自 crate 的 autodiff feature 后面，不开就没有（见 tenferro 卷的 feature 一节）。tenferro 源码只认 tag v0.8.0，链接一律是 https://github.com/tensor4all/tenferro-rs/blob/v0.8.0/...。JAX 和 PyTorch 一侧只写在虚拟环境里跑出来的行为，不链接它们的源码。

每章的结构

每个内容页的顺序固定：

1. 「主线位置」：这一页是哪条主线的哪一步。
2. 版本行和一张图。
3. 正文按对照点分节。每一节先给 JAX（或 PyTorch）代码，再给 tenferro 代码，然后是两边的真实输出。先说概念上相同的部分，再说哪里不同、为什么不同；tenferro 内部的机制链到 tenferro 卷。
4. 说「对得上」的地方，给出最大绝对误差和输出的量级，不是 0 时说明原因。
5. ## 小结。

一对对照程序

变换主线的主例子：f(x) = Σ sin(x)·x²，x 是 [−2, 2] 上等距的 1000 个点，求 ∇f。

JAX 一侧，grad_jax.py：

import sys

import jax

jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp
import numpy as np

N = 1000
X = np.array([-2 + 4 * i / (N - 1) for i in range(N)])


def f(x):
  return jnp.sum(jnp.sin(x) * x**2)


g = jax.jit(jax.grad(f))(jnp.array(X))
np.save(sys.argv[1], np.asarray(g))

tenferro 一侧，一个只依赖 tenferro-runtime、tenferro-cpu、tenferro-ad 的 crate：

use tenferro_ad::TracedTensorAdExt;
use tenferro_cpu::CpuEngine;
use tenferro_runtime::{DType, ExecutionEngine, GraphCompiler, Runtime, Tensor, TensorRead, TracedTensor};

const N: usize = 1000;

fn main() -> Result<(), Box<dyn std::error::Error>> {
  let xs: Vec<f64> = (0..N).map(|i| -2.0 + 4.0 * i as f64 / (N - 1) as f64).collect();
  let x = TracedTensor::input_concrete_shape(DType::F64, &[N])?;
  let y = (&x.sin()? * &(&x * &x)?)?.reduce_sum(None)?;
  let g = y.grad(&x)?;
  let program = GraphCompiler::new().compile_with_input_specs(&g, &[(&x, DType::F64, &[N][..])])?;

  let mut engine = CpuEngine::serial();
  let mut builder = Runtime::builder();
  builder.register_engine(tenferro_cpu::runtime_engine_registration(&engine)?)?;
  let runtime = builder.build()?;
  let xt = Tensor::from_vec_col_major(vec![N], xs)?;
  let reads = [TensorRead::from_tensor(&xt)];
  let prepared = runtime.prepare(&program, &reads, &[engine.target()])?;

  let mut out = runtime.run(&mut [&mut engine], &prepared, &reads)?;
  let grad = out.remove(0).into_tensor(&mut engine)?;
  let grad = grad.as_slice::<f64>()?;
  // ... 把 grad 按 {:.17e} 每行一个写进第一个命令行参数给的文件
  Ok(())
}

两段程序一步对一步：

这几步各自做什么、和 jit 的缓存有什么不同，见记图与 jit；compile 和 prepare 内部见 tenferro 卷的编译。

比较用一个几十行的 compare_vec.py：读两份向量，打印最大绝对误差、出现的位置、参考值的最大绝对值和逐位相同的个数；最大误差超过 tol × max(1, max|ref|) 就退出码 1。

ref, got = load(sys.argv[1]), load(sys.argv[2])
tol = float(sys.argv[3]) if len(sys.argv) > 3 else 1e-14
assert ref.shape == got.shape, (ref.shape, got.shape)
d = np.abs(ref - got)
scale = float(np.max(np.abs(ref)))
i = int(np.argmax(d))
bound = tol * max(1.0, scale)
print(f"n={ref.size} max|ref|={scale:.3e} max|diff|={d[i]:.3e} at {i} ({ref[i]!r} vs {got[i]!r}) "
      f"bitwise equal {int((ref == got).sum())}/{ref.size} bound={bound:.1e} {'OK' if d[i] <= bound else 'MISMATCH'}")
sys.exit(0 if d[i] <= bound else 1)

跑法：

python -I grad_jax.py grad_jax.npy
cargo run --release --offline -- grad_tf.txt
python -I compare_vec.py grad_jax.npy grad_tf.txt 1e-14

n=1000 max|ref|=3.153e+00 max|diff|=4.441e-16 at 6 (np.float64(2.0928852607508746) vs np.float64(2.092885260750874)) bitwise equal 761/1000 bound=3.2e-14 OK

退出码 0。一个从来不报错的检查等于没有检查，所以反过来验一次：把 tenferro 结果的第 500 个元素加 1e-13，再比：

n=1000 max|ref|=3.153e+00 max|diff|=1.000e-13 at 500 (np.float64(1.2024022661265687e-05) vs np.float64(1.2024022761265688e-05)) bitwise equal 761/1000 bound=3.2e-14 MISMATCH

退出码 1。

「对得上」是什么意思

看绝对误差，同时给出量级

1000 个梯度分量里最大的 |ref| 是 3.153，最大差 4.441e-16 出现在第 6 个元素，值约 2.09。[2, 4) 区间里 f64 的 ulp 正好是 2⁻⁵¹ ≈ 4.441e-16，所以这是 1 个 ulp：两边在这个元素上只差最后一位。1000 个里有 761 个逐位相同。hvp（jvp(grad)，切向量 vᵢ = cos i）也一样：最大 |ref| 5.148，最大差 1.332e-15，719 个逐位相同。

这一卷判断浮点结果用绝对误差，并且总写出最大 |ref|，让读者知道误差相对于什么量级。

为什么不是 0

两边算的是同一个数学式子，结果却不逐位相同。在这个例子里，239 个不同的元素有两个来源，都能单独验证。

一、sin 和 cos 不是同一个实现。 把 x 原样过一遍 jnp.sin（jit）、np.sin 和 tenferro 的 traced sin：

sin: jax   vs numpy    bitwise equal 1000/1000  max|d|=0.000e+00
sin: jax   vs tenferro bitwise equal  961/1000  max|d|=1.110e-16
sin: numpy vs tenferro bitwise equal  961/1000  max|d|=1.110e-16
cos: jax   vs numpy    bitwise equal 1000/1000  max|d|=0.000e+00
cos: jax   vs tenferro bitwise equal  946/1000  max|d|=1.110e-16
cos: numpy vs tenferro bitwise equal  946/1000  max|d|=1.110e-16

JAX 和 NumPy 在这 1000 个点上逐位相同；tenferro 的 sin 有 39 个、cos 有 54 个差 1 个 ulp。tenferro 的 sin、cos、exp 等在 simd_math::apply 里交给 elem-simd-math，一个 SLEEF 的 Rust 移植，不调平台的 libm（见 tenferro 卷的逐元素）。tenferro 的 sin 和 Rust 的 f64::sin 比，也是 961 个相同。

sin: elements where numpy/jax and tenferro differ: 39; there numpy/jax correctly rounded 12, tenferro correctly rounded 27; overall correctly rounded: numpy/jax 973/1000, tenferro 988/1000
cos: elements where numpy/jax and tenferro differ: 54; there numpy/jax correctly rounded 20, tenferro correctly rounded 34; overall correctly rounded: numpy/jax 962/1000, tenferro 976/1000

二、jit 把乘加融合成 FMA。 JAX 的梯度程序（jax.make_jaxpr(jax.grad(f))）是这样：

    b:f64[3] = sin a
    c:f64[3] = cos a
    d:f64[3] = integer_pow[y=2] a
    e:f64[3] = mul 2.0:f64[] a
    ...
    j:f64[3] = mul h e
    k:f64[3] = mul i c
    l:f64[3] = add_any j k

也就是 sin(x)·(2x) + x²·cos(x)。用 NumPy 按这个顺序逐步算一遍，和不加 jit 的 jax.grad(f) 1000 个全部逐位相同；和加了 jit 的只有 777 个相同。换成 Python 的 math.fma(2x, sin x, x²·cos x)（乘加只舍入一次），和 jit 的结果有 993 个相同：

no-jit grad == plain formula 1000 /1000; jit grad == plain 777 ; jit == fma(2x,s,x*x*c) 993 ; jit == fma(x*x,c,s*2x) 790 ; jit == either 994
no-jit vs jit max|d| 4.440892098500626e-16

XLA 在 CPU 上编译时把 mul 和 add_any 合成了一条 FMA，剩下 6 个元素用哪种结合方式本页没有追查。这说明 JAX 自己加不加 jit，结果就差 1 个 ulp。同一个融合在 SVD 求导规则里会造成大麻烦（对角线上本该是 0 的差变成 1e-16），见简并奇异值。

运算顺序是一样的。 把 tenferro 自己的 sin、cos 代进同一个式子 sin(x)·(2x) + x²·cos(x)，用 NumPy 逐步算，和 tenferro 的梯度 1000 个全部逐位相同：tenferro 的求导程序和 JAX 的 jaxpr 做的是同样的乘法和加法。把两个来源分开看：

jax no-jit vs tenferro: bitwise 956 /1000, max|d| 4.440892098500626e-16
jax jit vs jax no-jit: bitwise 777 /1000
jax jit vs tenferro: bitwise 761 /1000

不加 jit 的 JAX 和 tenferro 之间只差 sin、cos，44 个元素不同；JAX 加不加 jit 之间只差 FMA，223 个不同；两者叠在一起，就是 jit 的 JAX 和 tenferro 之间的 239 个。没有哪一边算错，两个正确的 f64 实现之间就是这个距离。

线程数也会改末位

上面的例子在 tenferro 里和线程数无关：tenferro 的 ad_vec 实验在单线程和 8 线程下 grad 逐位相同。求和、分解就不一样了。同样 100 万个数（xᵢ = ((7919·i) mod 1000003) / 1000003 − 0.5，两边都能逐位重建），先在 Python 里用四种方式求和：

fsum (exact)     = -1.95248614254157249e+00
sequential loop  = -1.95248614254587283e+00
numpy sum        = -1.95248614254195907e+00
jax jit sum      = -1.95248614254401787e+00

精确值（math.fsum）、逐个累加、NumPy、JAX 是四个不同的数。tenferro 一侧，同一组数在 CpuEngine::serial() 和 8 线程的 engine 上各算一次 reduce_sum，再对一个 256×256 矩阵做 svdvals：

sum serial = -1.95248614254350450e0  w8 = -1.95248614254227348e0
reduce_sum 1e6               n=    1 max|serial-w8|=1.231e-12 max|x|=1.952e0 bitwise equal 0/1
...
svdvals 256x256              n=  256 max|serial-w8|=2.043e-14 max|x|=4.110e1 bitwise equal 13/256
serial run twice bitwise equal: sum true matmul true svd true
w8 run twice bitwise equal: sum true matmul true svd true

单线程和 8 线程的和又是两个数，256 个奇异值只有 13 个逐位相同。但同一个宽度连跑两次，结果逐位相同。所以末位的差来自「宽度不同，工作的切法不同，加的顺序就不同」，而不是运行时的随机性。tenferro 不替你选宽度，宽度是建 Pool 时写死的（见 12 个线程比 8 个慢）。

TRG 也一样。20 步、χ = 16 的 lnZ/N 和 d(lnZ/N)/dβ，五个 β 上取最大差：

JAX (lorentz custom_jvp, jit) vs tenferro trg_staged_chi16_rank_w1.txt: max|dlnZ/N|=5.551e-16 max|dgrad|=2.424e-12
JAX (lorentz custom_jvp, jit) vs tenferro trg_staged_chi16_rank_w8.txt: max|dlnZ/N|=2.442e-15 max|dgrad|=4.887e-12
tenferro serial vs 8 threads: max|dlnZ/N|=1.998e-15 max|dgrad|=2.462e-12

lnZ/N 在 0.8 到 1.2 之间，差在 1e-15 量级；梯度在 0.7 到 1.9 之间，经过 40 次 SVD 的反向，差到了 1e-12 量级。关键是第三行：tenferro 自己只换线程数，差就有 2.0e-15 和 2.5e-12，和它与 JAX 之间的差是同一个量级。JAX 一侧用的 SVD 求导规则和 tenferro 默认的规则相同（custom_jvp 写的洛伦兹展宽），所以这里比的是同一个数学，不同的实现；JAX 内置的 jnp.linalg.svd 在这条链上梯度是 NaN（简并奇异值）。

容差

• 变换主线和其他小规模的 f64 结果：1e-14 × max(1, max|ref|)，相当于几十个 ulp。本页 grad 实测 4.4e-16，hvp 1.3e-15，留了 20 倍以上的余量；上面加的 1e-13 超出。
• TRG：lnZ/N 用 1e-14，导数用 1e-11。 导数的容差放宽，因为 tenferro 自己换线程数就差 2.5e-12；1e-11 留了 4 倍。
• 和精确解比是另一回事。 χ = 16 的截断误差在 lnZ/N 上是 7e-9 到 5e-6，在内能 u 上到 5.6e-4（β = 0.44，靠近临界点）。这是物理近似的误差，两边一样；要求两边彼此之间的差远小于它们各自和精确解的差。
• 整数、形状、dtype 要求完全相同，复数结果按约定换算（差一个共轭）之后再比（复数求导）。这些不涉及舍入，差了就是行为不同，比如整数除法的取整方向（整数运算）。

三条主线落在哪一页

TRG 主线（全卷的轴）：局部张量 → 一步 SVD 截断加收缩 → 累加 lnZ → 分段求导得内能 → 对照精确解。

变换主线：写函数 → 记图 → grad / vjp / jvp / hvp → 编译、准备 → 执行。

张量主线：建数组 → 计算 → 读回主机。

阅读路线

• 要搬一段 JAX 的张量网络代码：02 变换 → 05 线性代数 → 08 程序结构 → 10 动手，遇到数组上的意外再回 01 查。
• 只关心值不值得换：什么时候用 tenferro → 07 性能。
• 习惯 PyTorch 的 backward：03 eager，再看 06 复数的约定。
• 已经在写 tenferro、被 JAX 习惯绊住：查 09 常见坑。

动脑筋答案

4096 个全部逐位相同：

matmul 64x4096 * 4096x64     n= 4096 max|serial-w8|=0.000e0 max|x|=8.920e0 bitwise equal 4096/4096

• 矩阵乘也是求和，结果却和单线程一样。这说明在这个形状上，8 线程的 engine 没有改变每个输出的求和顺序：工作最可能是按输出切块分给线程，而不是把 4096 项的收缩维切开。本页没有去读 tcontract 和 faer 的分块代码来证实这一点。
• reduce_sum 只有一个输出，要并行就只能把 100 万个数切开，各自求部分和再相加，顺序随宽度变；SVD 内部也有这样的归约。
• 所以「换线程数会不会改末位」要按运算、按形状看，不能一概而论。本页只量了这一个形状，别的形状、别的宽度要自己量。和 JAX 对照时，tenferro 一侧总写明 engine 宽度。

小结

• 版本锁定 JAX 0.12.0.1（CPU，x64）、PyTorch 2.15.0、tenferro-rs 0.8.0（native）；tenferro 源码链接只认 tag v0.8.0。
• 每个对照是一个 .py 加一个 cargo crate，结果写成文件，compare_vec.py 报最大绝对误差和逐位相同的个数，超出容差退出码 1；把一个元素加 1e-13，它报 MISMATCH。
• f(x) = Σ sin(x)·x² 的梯度，两边最大差 4.4e-16，是 1 个 ulp，761 / 1000 逐位相同。来源有两个：tenferro 的 sin / cos 是 SLEEF 式实现而不是 libm，jit 把乘加融合成 FMA；两边求导程序的运算顺序相同。
• tenferro 换线程数，reduce_sum 和 SVD 的末位会变，同一宽度逐位可重复；TRG 上它自己单线程对 8 线程的差（2.0e-15 / 2.5e-12）和它对 JAX 的差是同一量级。
• 容差：小规模 f64 结果 1e-14 × max(1, max|ref|)；TRG 的 lnZ/N 1e-14、导数 1e-11；整数、形状、dtype 完全相同；和精确解的差是物理误差，另算。