整张图求导，几步之后就停不下来
主线位置：TRG 主线「分段求导得内能」：为什么要分段，而不是像 JAX 那样整张图
grad。
版本：JAX 0.12.0.1（CPU，x64）、tenferro-rs 0.8.0（CPU，native，单线程）
flowchart TB jx["jax.grad 整张图"] --> lin["变换随步数线性"] port["照搬 ad.grad 整张图"] --> reuse["值被反复使用"] reuse --> exp["变换随深度指数"] exp --> hang["5 步不返回"] port --> stage["按步分段"] stage --> vjp["每步一个 vjp"] vjp --> fast["20 步 0.56 s"]
实验代码：tenferro 一侧是 bin p25_whole_graph（release，CpuEngine::serial()）和动手页的分段程序 ising_trg，JAX 一侧是 p25_whole_graph.py。χ = 16，截断数 [2, 4, 4, 16, …]，β = 0.3。负载均值 2.4 到 3.0。M2 Max、macOS 27、rustc 1.99.0。
现象
JAX 里整张 TRG 图直接 jax.grad，20 步也没有问题：
jax.make_jaxpr(jax.grad(f))(x) # 追踪加变换
jax.jit(jax.grad(f))(x).block_until_ready() # 第一次调用，含编译
trg steps= 1 eqns 266 make_jaxpr(grad) 39.06 ms first jit(grad) call 104.86 ms grad=1.667309214024311
trg steps= 4 eqns 837 make_jaxpr(grad) 68.01 ms first jit(grad) call 252.37 ms grad=0.844054073623855
trg steps= 5 eqns 1038 make_jaxpr(grad) 68.52 ms first jit(grad) call 361.30 ms grad=0.744793433247367
trg steps=20 eqns 4053 make_jaxpr(grad) 211.57 ms first jit(grad) call 1104.63 ms grad=0.704498679529762
照搬到 tenferro，同样一张图，同样一次 ad.grad：
let beta = TracedTensor::input_concrete_shape(DType::F64, &[])?;
let lnz = trg_with(&beta, &schedule(16, n), &mut |_| Ok(()))?;
let g = ad.grad(&lnz, &beta)?; // 计时
// 编译、准备、执行 g；最后单独计时 drop(g)
trg n= 1 ad.grad 4.36 ms | drop 0.21 ms | grad ops 266 | run 17.867 ms | grad[0]=1.667309214024311
trg n= 2 ad.grad 5.27 ms | drop 3.63 ms | grad ops 465 | run 0.622 ms | grad[0]=1.277612005080206
trg n= 3 ad.grad 88.68 ms | drop 84.98 ms | grad ops 665 | run 1.835 ms | grad[0]=1.023802944795145
trg n= 4 ad.grad 2088.32 ms | drop 2093.47 ms | grad ops 863 | run 3.564 ms | grad[0]=0.844054073623855
trg: stopped after n=4 (ad.grad > 1500 ms cap)
单独让 5 步跑，timeout 60 到点终止（exit=124），一分钟里 ad.grad 没有返回：
20:25:28
...
trg n= 4 ad.grad 2092.07 ms | drop 2082.67 ms | grad ops 863 | run 2.267 ms | grad[0]=0.844054073623855
exit=124
20:26:28
- 3 步到 4 步变换乘了 24；释放导数图（
drop）又花一份同样的时间。 - 导数程序的算子数是线性的（每步约 200 个，1 步时和 JAX 的 266 个方程相同），执行只要几毫秒。慢的只是变换。
- 梯度在 1、2、4 步和 JAX 逐位相同，3 步差末位。
出处
不限于 SVD。最小的复现是 y ← (y·y)/2 的链，同一个值每层用两次：
mul n=16 ad.grad 4.13 ms | drop 4.00 ms | grad ops 81 | ...
mul n=20 ad.grad 62.46 ms | drop 62.08 ms | grad ops 101 | ...
mul n=24 ad.grad 1004.20 ms | drop 1004.57 ms | grad ops 121 | ...
mul n=25 ad.grad 2013.58 ms | drop 2003.01 ms | grad ops 126 | ...
每层正好乘 2；JAX 同样的链 30 层 make_jaxpr(grad) 7.0 ms。时间耗在 computegraph 的 ValueKey::eq 上：它按结构递归比较两个值的键，把 DAG 展开成树来走，底数就是值被使用的次数。栈采样、源码和来历见 求导变换的代价。TRG 的一步里 T 用了好几次、两次 SVD 的三个输出都用上，所以底数比 2 大得多。eager 的 backward 用同一个变换，整图 4 步每次约 7 秒（eager 还是 traced）。这是上游的性能问题，0.8.0 复现。
修法
不让一次变换看见很深的图：每一步单独建图、单独求 vjp，步与步之间手工把余切传回去，相当于在 JAX 里对每一步调 jax.vjp（分段，完整程序见 Ising 模型的 TRG）：
// 前向：每步在一个新叶子上建一张 eager 图
let leaf = rt.variable_from(cur.to_tensor(eng)?)?;
let (tn, ln) = rt.with_eager_session(eng, |s| step(s, &leaf, d, k))?;
// 反向：从最后一步往前，T_i 经 T_{i+1} 和 ln‖T_i‖ 两条路到 lnZ
let a = rt.vjp(eng, tn, leaf, &ct)?;
let b = rt.vjp(eng, ln, leaf, &c)?;
tenferro 0.8.0 eager staged chi=16 steps=20 width=1 h=1e-5 ks=[2, 4, 4, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16]
beta=0.3000 lnZ/N=0.790559063800966 u=-0.704498679529762 c=0.286288378175 | value_and_grad 555.8 ms
20 步的值加梯度 556 ms（第一次之后再跑 3 次取最小值；c 是另外两次求梯度的差分，不在其中）；u = −grad 和 JAX 整图的 0.704498679529762 打印出来相同。
怎么确认没中招
- 写好一步就量一次
ad.grad的时间，再加一步量一次：比值稳定在 2 以上就是指数增长，别等到整张图。 - 慢的程序先看是不是卡在变换或
drop上：执行只要毫秒，变换却要秒，就是这个坑。 - 用循环反复复用上一轮结果的图（迭代求解、多步演化）都要分段，不只 TRG。
小结
- tenferro 0.8.0 的
ad.grad在值被反复使用的深图上随深度指数增长，释放导数图同样慢；JAX 的grad是线性的。 - TRG 整图求导 3 步 88 ms、4 步 2.1 s、5 步一分钟不返回。
- 按步分段、逐步
vjp，20 步 556 ms。