从 JAX / PyTorch 到 tenferro

这一卷写给用 JAX 做数值计算、张量网络的读者，想把代码搬到 tenferro-rs 上。tenferro 的 traced 模式和 JAX 走的是同一条路：先记图，求导是对图的变换，产物还是图，编译、准备一次，再反复执行；它的 eager 模式是 PyTorch 式的，每个运算立即出数，标量上调 backward。所以对照的主角是 JAX，PyTorch 只在 eager 求导和复数求导约定这两处出场。

每章只做一件事：给出 JAX（或 PyTorch）和 tenferro 两份都能跑的程序，再贴两边的真实输出。先说两边概念上相同的部分，再说哪里不同、为什么不同。说「两边对得上」的地方，都给出最大绝对误差和量级，不是 0 时说明原因。

tenferro 内部怎么实现，这一卷不重讲：机制都在 tenferro-rs 源码解析那一卷，这里只给链接。Rust 语法也不讲。

三条主线

• TRG 主线（全卷的轴）：二维 Ising 模型的张量重整化群（TRG）。局部张量 → 一步 SVD 截断加收缩 → 累加 lnZ → 分段求导得内能 u = −∂(lnZ/N)/∂β → 对照 Onsager 精确解。键维数 χ = 16，20 步，β 取 0.30、0.40、0.44、0.50、0.60。
• 变换主线：写函数 → 记图（相当于 jax.make_jaxpr）→ grad / vjp / jvp / hvp → 编译、准备（相当于 jit 的第一次调用）→ 执行。例子是 f(x) = Σ sin(x)·x²，n = 1000。
• 张量主线：建数组 → 计算 → 读回主机。NumPy / JAX 数组对 tenferro 张量。

每个内容页开头的「主线位置」说明这一页在哪条线的哪一步。每一步落在哪一页，见怎么读这一卷。

flowchart TB
  intro[00 导读] --> tensor[01 数组]
  tensor --> trans[02 变换]
  trans --> eager[03 eager]
  trans --> ein[04 einsum]
  ein --> la[05 线性代数]
  la --> cplx[06 复数]
  la --> perf[07 性能]
  perf --> struct[08 程序结构]
  eager --> pit[09 常见坑]
  struct --> pit
  pit --> hands[10 动手]

实验环境

每个对照是一对程序：一个 .py 和一个 cargo crate，各自把结果写成文件，再用几十行的 Python 脚本逐个数比较。平台是 M2 Max、macOS 27、rustc 1.99.0。tenferro 一律离线构建；JAX 只用 CPU 设备。比速度的地方用 release 构建，JAX 先调一次（含编译）再计稳态，每次都 block_until_ready()；tenferro 写明 engine 的线程宽度；取多次运行的最小值，并写明当时机器的负载。

各部分

1. 00 导读：版本；一对对照程序怎么跑、f64 下按什么标准算对上，为什么不是 0；什么时候用 tenferro，什么时候留在 JAX。
2. 01 数组：列主序与 reshape 的顺序；dtype 的默认值、x64、类型提升；整数除法、除零、溢出。
3. 02 变换：记图对 jaxpr，占位符对 jit 的参数，prepare / run 对 jit 的缓存；grad、vjp、jvp、hvp；jvp(grad) 和 grad(grad)。
4. 03 eager：with_eager_session 和 backward 对 torch.autograd；eager 和 traced 怎么选。
5. 04 einsum：jnp.einsum / opt_einsum 的收缩路径对 tenferro 的规划；traced 张量上用 dot_general。
6. 05 线性代数：svd、eigh、qr、solve 的梯度；简并奇异值上 tenferro 默认的展宽规则对 JAX、PyTorch 的 NaN；SVD 输出的截断。
7. 06 复数：tenferro、PyTorch、JAX 三种复数求导约定并排。
8. 07 性能：计时方法；TRG 和向量函数的时间；求导变换本身的代价。
9. 08 程序结构：lax.scan 式的写法在 tenferro 里怎么落地：按形状分段的小程序、手工串反向；jax.checkpoint 对 checkpoint。
10. 09 常见坑：从 JAX 习惯出发的坑，每条有报错原文或现象和最小复现。
11. 10 动手：完整的 TRG，JAX 参考实现对 tenferro 分段实现，lnZ、u、比热 c 对照 Onsager。

适合谁读

• 会写 JAX（jit、grad、jvp、custom_jvp），能读懂 Rust 代码。
• 只想知道值不值得换：先读什么时候用 tenferro，那一页同时写了该留在 JAX 的情况。
• 手上有一段 JAX 的张量网络代码要搬：从怎么读这一卷看对照的标准，然后读 02 变换和 05 线性代数，最后看 10 动手里的完整 TRG。