程序结构

JAX 里一个 20 步的迭代算法可以直接写成 Python 循环，整个交给 jit(grad(f))；嫌编译慢就换 lax.scan，嫌内存大就套 jax.checkpoint。tenferro 里整图求导在深图上随深度指数增长（求导变换的代价），又没有 scan 算子，所以程序的结构要自己定：一步一个小程序，反向手工串。这一部分讲这种写法，以及两边的 checkpoint 各是什么。完整的 TRG 程序见 动手：Ising 模型的 TRG。

1. 分段：lax.scan 只扫形状相同的 16 步，编译时间减半；tenferro 按 (d, k) 建 5 种步的前向和反向程序，β、T、余切都是占位符，27 ms 准备好、所有 β 复用；反向用替身 ⟨T_{n+1}, ct_{n+1}⟩ + cₙ ln‖Tₙ‖ 逐步串，相当于逐步的 jax.vjp。
2. checkpoint：jax.checkpoint 重算换内存（197 → 35 MiB），在这条链上却让梯度偏 1e-4；tenferro 的 checkpoint 现在求值、旧图留给求导，4 步的 ad.grad 从 2.1 s 降到 4 ms，但检查点链表合并不去重，5 步栈溢出。tenferro 里对应 jax.checkpoint 的是 traced 分段。

flowchart TB
  loop["Python 循环 20 步"] --> scan["lax.scan"]
  loop --> remat["jax.checkpoint"]
  scan --> stg["tenferro 按形状分段"]
  remat --> stg
  stg --> tr["traced 反向重算"]
  stg --> eg["eager 反向用残差"]
  ck["tenferro checkpoint"] --> chain["链表每步乘 24"]