Ising 模型的 TRG
主线位置：TRG 主线全程，局部张量 → 一步 SVD 截断加收缩 → 累加 lnZ → 分段求导得内能 → 对照精确解，JAX 和 tenferro 各走一遍。
版本：JAX 0.12.0.1（CPU，x64）、tenferro-rs 0.8.0（CPU，native）
flowchart TB b["β"] --> t0["局部张量 T0"] t0 --> leaf["T_i 做成新叶子"] leaf --> stp["SVD 截断加收缩"] stp --> nxt["T_i+1 和 ln 范数"] nxt -->|"下一步"| leaf nxt --> lnz["ln tr 得 lnZ/N"] lnz --> rev["逆序逐步 vjp"] rev --> u["到 β 得 u"] u --> c["β±h 差分得 c"] c --> ons["对照 Onsager"]
前面各章各讲了 TRG 的一步。本页把它们拼成一个完整程序，JAX 一份、tenferro 一份，同一个算法、同一个收缩顺序，跑同样的 β，再和精确解逐项比。实验代码是 Python 脚本 ising_trg.py 和 tenferro crate 的 bin ising_trg（release），另有几个探查用的小脚本，下文用到时说明。M2 Max、macOS 27、rustc 1.99.0；同机有别的任务在跑，计时时 1 分钟负载均值 6–10。
要算什么
二维正方格子上的 Ising 模型，J = 1，无外场。每个格点的自由能、内能和比热：
- lnZ/N，N 是格点数；
- 内能 u = −∂(lnZ/N)/∂β；
- 比热 c = β² ∂²(lnZ/N)/∂β² = −β² ∂u/∂β。
Onsager 给出了 N → ∞ 的精确解，临界点 β_c = ln(1 + √2)/2 ≈ 0.440687；比热在 β_c 处对数发散。精确值用 SciPy 的椭圆积分算（onsager.py），这里只用它的输出。
TRG（Levin–Nave 的张量重整化群）把配分函数写成每个格点一个四脚张量 T[u, l, d, r] 的网络。每一步把 T 按两种对角方式拆成两个三脚张量（各做一次 SVD，只留前 k 个奇异值），再把四个三脚张量收缩成新的 T：张量数减半，键维数变成 k。每步先把 T 除以它的范数，范数的对数按这一步还剩的张量数累加进 lnZ。20 步之后只剩一个张量，取迹。这样算的是 N = 2²⁰ 个格点、周期边界的环面，lnZ/N 写成
lnZ/N = Σᵢ 2⁻ⁱ ln‖Tᵢ‖ + 2⁻²⁰ ln tr T₂₀，i = 0 … 19。
键维数上限 χ = 16。每步的截断数 k = min(D², χ)，前三步再压到 2、4、4：这三步的 D² 个奇异值里有精确为 0 的（结构性的秩），留着它们，sqrt(s) 的导数就是无穷大。tenferro 的 SVD 在这里给出精确的 0，梯度是 NaN（本页的 traced 分段程序换成不压秩的 k = [4, 16, 16, 16]，4 步就是 grad=NaN）；原因和对策见 SVD 的输出。所以两边都用 k = [2, 4, 4, 16, 16, …]。
JAX 参考实现
ising_trg.py 的主体：
EPS = 1e-12 # tenferro's default derivative_eps
def schedule(chi, steps):
ks, d = [], 2
for n in range(steps):
k = min(d * d, chi, [2, 4, 4][n]) if n < 3 else min(d * d, chi)
ks.append(k)
d = k
return ks
@jax.custom_jvp
def svd(a):
return jnp.linalg.svd(a, full_matrices=False)
@svd.defjvp
def svd_jvp(primals, tangents):
# tenferro's rule: F = d / (d^2 + eps^2), d = s_j^2 - s_i^2, diagonal masked
(a,), (da,) = primals, tangents
u, s, vh = svd(a)
v = vh.T
dp = u.T @ da @ v
d = (s[None, :] - s[:, None]) * (s[None, :] + s[:, None])
f = jnp.where(jnp.eye(s.shape[0]) > 0, 0.0, d / (d * d + EPS * EPS))
dss, sds = dp * s[None, :], s[:, None] * dp
return (u, s, vh), (u @ (f * (dss + dss.T)), jnp.diagonal(dp), (v @ (f * (sds + sds.T))).T)
def initial_tensor(beta):
c, s = jnp.sqrt(jnp.cosh(beta)), jnp.sqrt(jnp.sinh(beta))
q = jnp.stack([jnp.stack([c, s]), jnp.stack([c, -s])])
return jnp.einsum("su,sl,sd,sr->uldr", q, q, q, q)
def make_trg(ks):
n = len(ks)
def trg(beta):
t = initial_tensor(beta)
acc = 0.0
for i, k in enumerate(ks):
norm = jnp.sqrt(jnp.sum(t * t))
t = t / norm
acc = acc + 2.0 ** (n - i) * jnp.log(norm)
d = t.shape[0]
ua, sa, vah = svd(t.transpose(2, 1, 0, 3).reshape(d * d, d * d, order="F"))
ub, sb, vbh = svd(t.transpose(3, 2, 1, 0).reshape(d * d, d * d, order="F"))
s1 = (ua[:, :k] * jnp.sqrt(sa[:k])).reshape(d, d, k, order="F")
s3 = (vah[:k, :].T * jnp.sqrt(sa[:k])).reshape(d, d, k, order="F")
s2 = (ub[:, :k] * jnp.sqrt(sb[:k])).reshape(d, d, k, order="F")
s4 = (vbh[:k, :].T * jnp.sqrt(sb[:k])).reshape(d, d, k, order="F")
t = jnp.einsum("war,abu,bgl,gwd->ruld", s1, s2, s3, s4)
acc = acc + jnp.log(jnp.einsum("xyxy->", t))
return acc / 2.0 ** n
return trg
四处值得说明，前两处各有一页讲原因：
- 自己写的 SVD 求导规则。 内置的
jnp.linalg.svd在这条链上梯度是 NaN（试验阶段实测从第 5 步起，原因是简并奇异值）。custom_jvp里写的是 tenferro 默认的规则：F = d/(d² + ε²)，ε = 1e-12，屏蔽对角线；(s_j − s_i)(s_j + s_i)而不是s_j**2 − s_i**2，是为了躲开jit的 FMA。见 简并奇异值。 reshape(..., order="F")。 tenferro 的 reshape 是列主序，JAX 这边改成 Fortran 序，两边的矩阵化逐位相同，SVD 的输入就一样（布局与 reshape）。不改也能算对，只是 SVD 的符号约定不同，末位跟着变。- Python 的
for把 20 步展开。jit追踪出一张 20 步的图，grad对整张图求导。截断数 k 是 Python 整数，是图的形状的一部分。 - 一步的收缩写成一个四操作数的 einsum。
jnp.einsum按 opt_einsum 的路径收缩，jnp.einsum_path给出的是先 S1·S2、再 S3·S4、最后两者相收缩；tenferro 的 einsum 规划对这个下标串选的是同一棵树（einsum）。写成三个两两相乘的链 ((S1·S2)·S3)·S4 也对，只是慢，见后面的时间表。
main 里求值和求导：
def main():
chi = int(sys.argv[1]) if len(sys.argv) > 1 else 16
steps = int(sys.argv[2]) if len(sys.argv) > 2 else 20
h = float(sys.argv[3]) if len(sys.argv) > 3 else 1e-5
betas = [float(x) for x in sys.argv[4:]] or [0.3, 0.4, 0.44, 0.5, 0.6]
ks = schedule(chi, steps)
print(f"jax {jax.__version__} x64={jax.config.jax_enable_x64} chi={chi} steps={steps} h={h:g} ks={ks}")
f = make_trg(ks)
vg = jax.jit(jax.value_and_grad(f))
g = jax.jit(jax.grad(f))
c2 = jax.jit(lambda b: jax.jvp(jax.grad(f), (b,), (jnp.ones_like(b),))[1])
for b in betas:
x = jnp.float64(b)
t = time.perf_counter()
lnz, gr = jax.block_until_ready(vg(x))
t_first = time.perf_counter() - t
t_best = float("inf")
for _ in range(5):
t = time.perf_counter()
jax.block_until_ready(vg(x))
t_best = min(t_best, time.perf_counter() - t)
lnz, gr = float(lnz), float(gr)
c_fd = b * b * (float(g(x + h)) - float(g(x - h))) / (2 * h)
c_jvp = b * b * float(c2(x))
print(f"beta={b:.4f} lnZ/N={lnz:.15f} u={-gr:.15f} c={c_fd:.12f} c_jvp={c_jvp:.12f} "
f"| value_and_grad first {t_first * 1e3:.1f} ms steady {t_best * 1e3:.1f} ms")
print(f" hex lnZ/N={lnz.hex()} grad={gr.hex()} c={c_fd.hex()}")
c_jvp 是 jvp(grad) 直接给出的二阶导数；c 是对 grad 做中心差分，和 tenferro 一侧的做法相同，用来逐位对照。
tenferro 实现
tenferro 不能照抄「一张图、整体求导」：对 20 步的整图求导，变换本身的耗时随步数指数增长，试验阶段实测 2 步 traced grad 7.7 ms，6 步超过 3 分钟（求导的代价）。办法是分段：每一步是一张独立的小图，反向由程序自己从最后一步往前串（分段）。分段有 traced 和 eager 两种写法，本页的主程序用 eager，理由在后面换成 traced 分段一节给出实测。
完整的 ising_trg.rs：
// 2D Ising model, Levin-Nave TRG in tenferro eager mode, staged by RG step.
// lnZ/N, u = -d(lnZ/N)/dbeta by reverse mode, c = beta^2 d^2(lnZ/N)/dbeta^2 as the central
// difference of the AD derivative. Same steps and contraction order as py/ising_trg.py.
// Reverse sweep: one EagerRuntime::vjp per step and output, like jax.vjp inside a loop.
// usage: ising_trg [chi=16] [steps=20] [width=1] [h=1e-5] [beta ...]
use std::time::Instant;
use tenferro_ad::{AdContext, EagerRuntime, EagerSession, EagerTensor, Tensor};
use tenferro_cpu::{CpuEngine, Pool};
use tenferro_einsum::EagerSessionEinsumExt;
use tenferro_linalg::EagerSessionLinalgExt;
use tenferro_runtime::SliceConfig;
type R<T> = Result<T, Box<dyn std::error::Error>>;
type Rt = std::sync::Arc<EagerRuntime>;
/// k per step: min(D^2, chi), and the first three capped at the structural rank 2, 4, 4.
fn schedule(chi: usize, steps: usize) -> Vec<usize> {
let mut ks = Vec::new();
let mut d = 2;
for n in 0..steps {
let k = if n < 3 { (d * d).min(chi).min([2, 4, 4][n]) } else { (d * d).min(chi) };
ks.push(k);
d = k;
}
ks
}
/// T[u,l,d,r] = sum_s q[s,u] q[s,l] q[s,d] q[s,r], q = [[c, s], [c, -s]].
fn initial(s: &mut EagerSession<'_>, beta: &EagerTensor, a: &EagerTensor, b: &EagerTensor) -> R<EagerTensor> {
let (e, en) = (s.exp(beta)?, s.neg(beta)?);
let en = s.exp(&en)?;
let ch = s.add(&e, &en)?;
let ch = s.scale_real(&ch, 0.5)?;
let sh = s.sub(&e, &en)?;
let sh = s.scale_real(&sh, 0.5)?;
let (c, sq) = (s.sqrt(&ch)?, s.sqrt(&sh)?);
let cb = s.broadcast_in_dim(&c, &[2, 2], &[])?;
let sb = s.broadcast_in_dim(&sq, &[2, 2], &[])?;
let qa = s.mul(&cb, a)?;
let qb = s.mul(&sb, b)?;
let q = s.add(&qa, &qb)?;
Ok(s.einsum(&[&q, &q, &q, &q], "su,sl,sd,sr->uldr")?)
}
/// One RG step: (T', ln ||T||), T' built from T / ||T||.
fn step(s: &mut EagerSession<'_>, t: &EagerTensor, d: usize, k: usize) -> R<(EagerTensor, EagerTensor)> {
let tt = s.mul(t, t)?;
let nn = s.reduce_sum(&tt, None)?;
let norm = s.sqrt(&nn)?;
let nb = s.broadcast_in_dim(&norm, &[d, d, d, d], &[])?;
let t = s.div(t, &nb)?;
let ln = s.log(&norm)?;
let m = d * d;
let ta = s.transpose(&t, &[2, 1, 0, 3])?;
let ma = s.reshape(&ta, [m, m])?;
let tb = s.transpose(&t, &[3, 2, 1, 0])?;
let mb = s.reshape(&tb, [m, m])?;
let split = |s: &mut EagerSession<'_>, mat: &EagerTensor| -> R<(EagerTensor, EagerTensor)> {
let (u, sv, vt) = s.svd(mat)?;
// reshape to the same concrete shape before slicing (05-linalg/svd-outputs)
let (u, sv, vt) = (s.reshape(&u, [m, m])?, s.reshape(&sv, [m])?, s.reshape(&vt, [m, m])?);
let sk = s.slice(&sv, SliceConfig { starts: vec![0], limits: vec![k], strides: vec![1] })?;
let sk = s.sqrt(&sk)?;
let skb = s.broadcast_in_dim(&sk, &[m, k], &[1])?;
let cols = SliceConfig { starts: vec![0, 0], limits: vec![m, k], strides: vec![1, 1] };
let uk = s.slice(&u, cols.clone())?;
let left = s.mul(&uk, &skb)?;
let v = s.transpose(&vt, &[1, 0])?;
let vk = s.slice(&v, cols)?;
let right = s.mul(&vk, &skb)?;
Ok((s.reshape(&left, [d, d, k])?, s.reshape(&right, [d, d, k])?))
};
let (s1, s3) = split(s, &ma)?;
let (s2, s4) = split(s, &mb)?;
let tn = s.einsum(&[&s1, &s2, &s3, &s4], "war,abu,bgl,gwd->ruld")?;
Ok((tn, ln))
}
/// ln sum_{x,y} T[x,y,x,y]
fn log_trace(s: &mut EagerSession<'_>, t: &EagerTensor) -> R<EagerTensor> {
let tr = s.einsum(&[t], "xyxy->")?;
Ok(s.log(&tr)?)
}
fn first(t: &EagerTensor) -> R<f64> {
Ok(t.value()?.as_slice::<f64>()?[0])
}
/// (lnZ/N, d(lnZ/N)/dbeta)
fn value_and_grad(rt: &Rt, eng: &mut CpuEngine, b: f64, ks: &[usize]) -> R<(f64, f64)> {
let n = ks.len();
let scale = 0.5f64.powi(n as i32);
let qa = rt.constant_from(Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 1.0, 0.0, 0.0])?)?;
let qb = rt.constant_from(Tensor::from_vec_col_major(vec![2, 2], vec![0.0_f64, 0.0, 1.0, -1.0])?)?;
let beta = rt.variable_from(Tensor::from_vec_col_major(vec![], vec![b])?)?;
// forward: each RG step is its own eager graph on a fresh leaf
let t0 = rt.with_eager_session(eng, |s| initial(s, &beta, &qa, &qb))?;
let (mut cur, mut d, mut acc) = (t0.clone(), 2, 0.0);
let mut stages = Vec::with_capacity(n); // (leaf T_i, T_{i+1}, ln ||T_i||)
for (i, &k) in ks.iter().enumerate() {
let leaf = rt.variable_from(cur.to_tensor(eng)?)?;
let (tn, ln) = rt.with_eager_session(eng, |s| step(s, &leaf, d, k))?;
acc += 2f64.powi((n - i) as i32) * first(&ln)?;
cur = tn.clone();
stages.push((leaf, tn, ln));
d = k;
}
let last = rt.variable_from(cur.to_tensor(eng)?)?;
let lt = rt.with_eager_session(eng, |s| log_trace(s, &last))?;
let lnz = (acc + first(<)?) * scale;
// reverse: ct = d(lnZ/N)/dT_i, newest step first, like jax.vjp per step
let seed = rt.constant_from(Tensor::from_vec_col_major(vec![], vec![scale])?)?;
let mut ct = rt.constant_from(rt.vjp(eng, <, &last, &seed)?.to_tensor(eng)?)?;
for (i, (leaf, tn, ln)) in stages.iter().enumerate().rev() {
let c = rt.constant_from(Tensor::from_vec_col_major(vec![], vec![2f64.powi((n - i) as i32) * scale])?)?;
// T_i reaches lnZ through T_{i+1} and through ln ||T_i||
let a = rt.vjp(eng, tn, leaf, &ct)?;
let b = rt.vjp(eng, ln, leaf, &c)?;
let sum = rt.with_eager_session(eng, |s| s.add(&a, &b))?;
ct = rt.constant_from(sum.to_tensor(eng)?)?;
}
let g = rt.vjp(eng, &t0, &beta, &ct)?.value()?.as_slice::<f64>()?[0];
Ok((lnz, g))
}
fn main() -> R<()> {
let args: Vec<String> = std::env::args().collect();
let chi: usize = args.get(1).map(|s| s.parse()).transpose()?.unwrap_or(16);
let steps: usize = args.get(2).map(|s| s.parse()).transpose()?.unwrap_or(20);
let width: usize = args.get(3).map(|s| s.parse()).transpose()?.unwrap_or(1);
let h: f64 = args.get(4).map(|s| s.parse()).transpose()?.unwrap_or(1e-5);
let betas: Vec<f64> = if args.len() > 5 {
args[5..].iter().map(|s| s.parse()).collect::<Result<_, _>>()?
} else {
vec![0.3, 0.4, 0.44, 0.5, 0.6]
};
let ks = schedule(chi, steps);
println!("tenferro 0.8.0 eager staged chi={chi} steps={steps} width={width} h={h:e} ks={ks:?}");
let mut rules = tenferro_linalg::semantic_ad_rules()?;
rules.merge(tenferro_einsum::semantic_ad_rules()?)?;
let ad = AdContext::builder().with_semantic_extension_rules(rules)?.build()?;
let rt = EagerRuntime::with_ad_context(&ad)?;
let mut eng = if width == 1 { CpuEngine::serial() } else { CpuEngine::new(Pool::builder().threads(width).build()?) };
for b in betas {
let (lnz, g) = value_and_grad(&rt, &mut eng, b, &ks)?;
let mut t_one = f64::INFINITY; // min of 3 more runs
for _ in 0..3 {
let t = Instant::now();
value_and_grad(&rt, &mut eng, b, &ks)?;
t_one = t_one.min(t.elapsed().as_secs_f64() * 1e3);
}
let (_, gp) = value_and_grad(&rt, &mut eng, b + h, &ks)?;
let (_, gm) = value_and_grad(&rt, &mut eng, b - h, &ks)?;
let c = b * b * (gp - gm) / (2.0 * h);
println!("beta={b:.4} lnZ/N={lnz:.15} u={:.15} c={c:.12} | value_and_grad {t_one:.1} ms", -g);
println!(" hex lnZ/N={:016x} grad={:016x} c={:016x}", lnz.to_bits(), g.to_bits(), c.to_bits());
}
Ok(())
}
和 JAX 一步对一步：
| JAX | tenferro | 说明 |
|---|---|---|
jnp.einsum("su,sl,sd,sr->uldr", q, q, q, q)、jnp.einsum("war,abu,bgl,gwd->ruld", ...) | s.einsum(&[&q, &q, &q, &q], "su,sl,sd,sr->uldr")、s.einsum(&[&s1, &s2, &s3, &s4], "war,abu,bgl,gwd->ruld") | 下标串相同；eager 会话上有 einsum，traced 张量上没有 |
t.transpose(2, 1, 0, 3).reshape(d * d, d * d, order="F") | s.transpose(&t, &[2, 1, 0, 3])、s.reshape(&ta, [m, m]) | 列主序，等于 order="F" |
svd(...)（custom_jvp） | s.svd(mat) | 默认规则就是展宽的那一条，不用自己写 |
ua[:, :k]、sa[:k] | 先 reshape 成同样的形状，再 s.slice | SVD 输出的长度是 Min(m, m)，直接切片求导会报错（SVD 的输出） |
jnp.log(jnp.einsum("xyxy->", t)) | s.einsum(&[t], "xyxy->")、s.log(&tr) | 重复下标取迹，两边都支持 |
整个 trg 是一个函数，jit(value_and_grad) | value_and_grad：每步 variable_from 一个新叶子，反向逐步 rt.vjp | 相当于每步 jax.vjp(step, t)，手工串起 pullback |
jax.config.update("jax_enable_x64", True) | 字面量 1.0_f64 | dtype 来自 Rust 字面量 |
前向一段：variable_from 用上一步的数值做一个新的叶子，with_eager_session 里的运算立即执行，同时记下这一步的 trace 和残差（两次 SVD 的结果）。叶子把 trace 切断：第 i 步的图只从 Tᵢ 走到 Tᵢ₊₁ 和 ln‖Tᵢ‖，不往前延伸。
反向一段：Tᵢ 通过两条路影响 lnZ/N，一条经过 Tᵢ₊₁，一条经过 ln‖Tᵢ‖ 的系数 2^(n−i)/2ⁿ。所以每步两次 vjp，余切分别是后一步传回来的 ct 和这个系数，相加就是 Tᵢ 的余切，再交给前一步。最后一次 vjp 从 T₀ 回到 β，得到 ∂(lnZ/N)/∂β。vjp 是函数式接口，返回普通的 EagerTensor，不碰叶子的梯度槽（eager AD 的「函数式的 grad、vjp、jvp」一节），所以同一个 EagerRuntime 反复调用 value_and_grad 不会把梯度累加起来。
📝NOTE
冷知识：eager 的 vjp 是 v0.3.0 才有的
PR #1304（「Implement eager AD end-state APIs」，2026-07-05 合并，首次发布于 v0.3.0）之前，EagerTensor 上只有一个不带参数的 backward() 和读槽位的 grad()：反向只能从一个标量出发，结果写进叶子的槽。这个 PR 在 EagerRuntime 上加了 grad、vjp、jvp（及 _optional 版本）和 no_grad，在 EagerTensor 上加了带余切的 backward_with。同一个 PR 改的 docs/spec/ad-contract.md 写明了两类接口的分工：「Stateful eager reverse mode (backward() and backward_with(seed)) accumulates reachable tracked leaves into gradient slots」，「Functional eager transforms (grad, vjp, jvp) return ordinary eager tensors and do not mutate gradient slots」。没有 vjp 时，上面的反向只能写成标量替身 ⟨Tᵢ₊₁, ct⟩ + c·ln‖Tᵢ‖ 再 backward，eager 还是 traced 用的就是那种写法；两种写法的结果见下文。
运行结果
python -I ising_trg.py
cargo run --release --offline --bin ising_trg -- 16 20 1
cargo run --release --offline --bin ising_trg -- 16 20 8
JAX：
jax 0.12.0.1 x64=True chi=16 steps=20 h=1e-05 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.704498679529761 c=0.286288378290 c_jvp=0.286288379149 | value_and_grad first 1233.9 ms steady 188.2 ms
hex lnZ/N=0x1.94c4285929bddp-1 grad=0x1.68b40d094f9edp-1 c=0x1.2528c7d7eb9dfp-2
beta=0.4000 lnZ/N=0.879363060306853 u=-1.106045407610774 c=0.861447339449 c_jvp=0.861447343854 | value_and_grad first 203.6 ms steady 199.7 ms
hex lnZ/N=0x1.c23be002a8255p-1 grad=0x1.1b25cab5944aap+0 c=0x1.b90fa02c52d01p-1
beta=0.4400 lnZ/N=0.928723962577071 u=-1.402785274913315 c=3.270066566479 c_jvp=3.311956762689 | value_and_grad first 222.3 ms steady 225.6 ms
hex lnZ/N=0x1.db81b50c8f629p-1 grad=0x1.671cef8f10249p+0 c=0x1.a2918a8f62a2fp+1
beta=0.5000 lnZ/N=1.025792858455947 u=-1.745583468150079 c=0.724737280977 c_jvp=0.724737277976 | value_and_grad first 196.0 ms steady 195.1 ms
hex lnZ/N=0x1.069a5c5b89e51p+0 grad=0x1.bede8ee42491bp+0 c=0x1.7310c3cffa2dfp-1
beta=0.6000 lnZ/N=1.210133028164416 u=-1.909086939792357 c=0.313435729607 c_jvp=0.313435729647 | value_and_grad first 191.2 ms steady 190.0 ms
hex lnZ/N=0x1.35cb4733c68ebp+0 grad=0x1.e8b9ebf3a0ff5p+0 c=0x1.40f54bc03cbbfp-2
tenferro，1 线程（CpuEngine::serial()）：
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 539.6 ms
hex lnZ/N=3fe94c4285929bdc grad=3fe68b40d094f9f3 c=3fd2528c7d5f23ef
beta=0.4000 lnZ/N=0.879363060306852 u=-1.106045407611296 c=0.861447342402 | value_and_grad 546.3 ms
hex lnZ/N=3fec23be002a824e grad=3ff1b25cab594dda c=3feb90fa045af001
beta=0.4400 lnZ/N=0.928723962577072 u=-1.402785274913331 c=3.270066566430 | value_and_grad 558.1 ms
hex lnZ/N=3fedb81b50c8f62c grad=3ff671cef8f10292 c=400a2918a8f47758
beta=0.5000 lnZ/N=1.025792858455948 u=-1.745583468150080 c=0.724737280960 | value_and_grad 543.7 ms
hex lnZ/N=3ff069a5c5b89e52 grad=3ffbede8ee42491e c=3fe7310c3cfd58ef
beta=0.6000 lnZ/N=1.210133028164416 u=-1.909086939792358 c=0.313435729623 | value_and_grad 540.7 ms
hex lnZ/N=3ff35cb4733c68ec grad=3ffe8b9ebf3a0ff7 c=3fd40f54bc0830c0
tenferro，8 线程：
tenferro 0.8.0 eager staged chi=16 steps=20 width=8 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.704498679529763 c=0.286288378181 | value_and_grad 269.3 ms
hex lnZ/N=3fe94c4285929bdc grad=3fe68b40d094f9fd c=3fd2528c7d60c9d0
beta=0.4000 lnZ/N=0.879363060306851 u=-1.106045407613180 c=0.861447337606 | value_and_grad 268.7 ms
hex lnZ/N=3fec23be002a8246 grad=3ff1b25cab596efd c=3feb90fa01c7c201
beta=0.4400 lnZ/N=0.928723962577072 u=-1.402785274913331 c=3.270066566421 | value_and_grad 275.3 ms
hex lnZ/N=3fedb81b50c8f62c grad=3ff671cef8f10291 c=400a2918a8f42bb7
beta=0.5000 lnZ/N=1.025792858455948 u=-1.745583468150080 c=0.724737280960 | value_and_grad 268.5 ms
hex lnZ/N=3ff069a5c5b89e53 grad=3ffbede8ee42491d c=3fe7310c3cfd58ef
beta=0.6000 lnZ/N=1.210133028164416 u=-1.909086939792358 c=0.313435729623 | value_and_grad 266.2 ms
hex lnZ/N=3ff35cb4733c68ec grad=3ffe8b9ebf3a0ff9 c=3fd40f54bc0830c0
逐位比较（compare.py 读两边的十六进制位模式，再算和 Onsager 精确解的差）：
tf_20_w1.txt vs jax_20.txt: max|d lnZ/N|=7.772e-16 max|ref|=1.210e+00 bitwise 0/5
tf_20_w1.txt vs jax_20.txt: max|d grad|=5.222e-13 max|ref|=1.909e+00 bitwise 0/5
tf_20_w1.txt vs jax_20.txt: max|d c(fd)|=2.952e-09 max|ref|=3.270e+00 bitwise 0/5
tf_20_w8.txt vs jax_20.txt: max|d lnZ/N|=1.665e-15 max|ref|=1.210e+00 bitwise 0/5
tf_20_w8.txt vs jax_20.txt: max|d grad|=2.406e-12 max|ref|=1.909e+00 bitwise 0/5
tf_20_w8.txt vs jax_20.txt: max|d c(fd)|=1.844e-09 max|ref|=3.270e+00 bitwise 0/5
beta lnZ/N err u err c err (JAX) c err (tenferro)
0.3000 -7.150e-09 +3.913e-07 -1.825e-06 -1.825e-06
0.4000 -7.605e-07 +3.380e-05 -2.510e-04 -2.510e-04
0.4400 -4.578e-06 -5.583e-04 +3.804e-01 +3.804e-01
0.5000 +4.576e-08 -1.889e-05 -1.342e-04 -1.342e-04
0.6000 +6.399e-07 -7.621e-07 -9.629e-06 -9.629e-06
- 两边彼此对得上。 lnZ/N 约是 1，两边差 7.8e-16；导数约是 0.7–1.9，差 5.2e-13。tenferro 换到 8 线程，和 JAX 的差变成 1.7e-15 和 2.4e-12：线程数改变了 SVD 和求和内部的加法顺序（怎么读这一卷 的「线程数也会改末位」）。
- c 两边差 3.0e-9，比导数的差大四个量级。 c 是导数的差分除以 2h = 2e-5，导数末位的差被放大了 5 万倍。
- 和精确解的差两边一样，到所列的位数相同。 这是 χ = 16 截断的物理误差，不是哪一边算错：lnZ/N 在 7e-9 到 5e-6，u 在 8e-7 到 6e-4，c 在 2e-6 到 0.38（β = 0.44 这一项还叠加了下一节讲的差分问题）。离临界点越近误差越大，β = 0.44 最差。
比热：差分步长要避开跳变
β = 0.44 那一行，c 的误差是 +0.38，而 JAX 一侧的 c_jvp 是 3.3120，c（h = 1e-5 的差分）是 3.2701，两个本该相同的数差了 1.3%。别的 β 上两者差 1e-9 量级。
先在 JAX 一侧换几个 h（smooth_044.py：对 jit 的求值和梯度做中心差分）：
beta=0.3 grad=0.704498679530 jvp(grad)=3.180981991
h=1e-02 fd(value)=0.704681004754 fd(grad)=3.182713683
h=1e-03 fd(value)=0.704500496691 fd(grad)=3.180998955
h=1e-04 fd(value)=0.704498697701 fd(grad)=3.180982148
h=1e-05 fd(value)=0.704498679693 fd(grad)=3.180981981
h=1e-06 fd(value)=0.704498679838 fd(grad)=3.180981980
h=1e-07 fd(value)=0.704498677173 fd(grad)=3.180981942
beta=0.44 grad=1.402785274913 jvp(grad)=17.107214683
h=1e-02 fd(value)=1.406263373179 fd(grad)=10.644008193
h=1e-03 fd(value)=1.403628368656 fd(grad)=16.241256587
h=1e-04 fd(value)=1.402814151670 fd(grad)=17.060222254
h=1e-05 fd(value)=1.402968993797 fd(grad)=16.890839703
h=1e-06 fd(value)=1.402785276106 fd(grad)=17.107214694
h=1e-07 fd(value)=1.402785277382 fd(grad)=17.107214568
β = 0.30 上，h 从 1e-2 到 1e-6，差分收敛到 AD 的值，再小就被舍入淹没。β = 0.44 上不单调：h = 1e-5 时对求值的差分偏了 1.8e-4，h = 1e-6 时又和 AD 一致到 1e-9。偏差只在 h 跨过某个点时出现，说明 lnZ/N(β) 在 0.44 附近有一个跳变。
扫一遍 [0.44 − 1e-5, 0.44 + 1e-5]，把每一小段上 lnZ/N 的增量和 AD 斜率预测的增量比，再二分（jump_044.py；最后两行用 NumPy 按同样的步骤重放一遍，打印最后一步截断处的第 16、17 个奇异值）：
jump in [0.440002500, 0.440003000]: dv - g*db = +3.701e-09
bisected: beta* in [0.440002964563000, 0.440002964563001], v jump +1.539e-09
below step 19: sa[k-1..k+1]=4.794909260534e-02 4.794909260532e-02 sb[k-1..k+1]=4.794909260534e-02 4.794909260532e-02
above step 19: sa[k-1..k+1]=4.794909260536e-02 4.794909260533e-02 sb[k-1..k+1]=4.794909260536e-02 4.794909260533e-02
第一行：在 [0.4400025, 0.4400030] 这一小段上，lnZ/N 比 AD 斜率预测的多涨了 3.7e-9，别的段上都在 1e-12 以内。第二行：二分到 β ≈ 0.440002964563，相邻两个 f64 的 β 上 lnZ/N 差 1.5e-9，是一个真正的跳变，不是陡而连续的变化。后两行：跳变两侧，最后一步截断处的第 16 和第 17 个奇异值（两个矩阵 sa、sb 都是）相同到 13 位有效数字。
再看最后一步（第 19 步）截断处的两个奇异值随 β 怎么走（cut_gap.py，只打印第 14 步以后，以及相对间隙小于 1e-6 的步）：
beta=0.44
...
step 18: s16=5.4430578888527e-02 s17=4.9506737627838e-02 (s16-s17)/s16=9.0e-02
step 19: s16=4.7948040886043e-02 s17=4.7740612264277e-02 (s16-s17)/s16=4.3e-03
beta=0.4400025
...
step 18: s16=5.4481531915543e-02 s17=4.9563226866255e-02 (s16-s17)/s16=9.0e-02
step 19: s16=4.7948961711535e-02 s17=4.7916363146855e-02 (s16-s17)/s16=6.8e-04
beta=0.440003
...
step 18: s16=5.4491734633316e-02 s17=4.9574465523668e-02 (s16-s17)/s16=9.0e-02
step 19: s16=4.7951590125494e-02 s17=4.7949102069628e-02 (s16-s17)/s16=5.2e-05
第 16 个和第 17 个奇异值在 β = 0.44 时相对差 4.3e-3，到 0.4400025 是 6.8e-4，到 0.440003 是 5.2e-5，在 β ≈ 0.44000296 交叉。交叉时，截断留下的第 16 个奇异向量从一个换成另一个，lnZ/N 跳了一下。这是截断 TRG 本身的性质，和用哪个框架无关：tenferro 的 h = 1e-5 差分给出同样的 3.27007，换成 h = 1e-6：
tenferro 0.8.0 eager staged chi=16 steps=20 width=1 h=1e-6 ks=[2, 4, 4, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16]
beta=0.4400 lnZ/N=0.928723962577072 u=-1.402785274913331 c=3.311956765457 | value_and_grad 557.0 ms
hex lnZ/N=3fedb81b50c8f62c grad=3ff671cef8f10292 c=400a7ee3304b3c30
c = 3.311956765，和 JAX 的 jvp(grad)（3.311956763）一致到 2.8e-9。
AD 给出的是跳变两侧各自光滑那一段上的导数，不受影响；受影响的是跨过跳变的差分。实用的做法：对 AD 能量做差分时用两个 h（比如 1e-5 和 1e-6）各算一次，对不上就说明中间有交叉，取小的那个，或者换一个 β。h 也不能太小：上表 h = 1e-7 时差分已经开始被导数的末位噪声主导。tenferro 一侧没有 jvp(grad) 可以对照，因为对整张图求二阶导数同样指数增长（高阶导数），而分段的二阶导数要把每步的 vjp 再做 jvp，本页没有写。
临界点附近扫一遍
h = 1e-6，β 从 0.40 到 0.48，JAX（c_jvp）和 tenferro（1 线程，对 AD 能量差分）各算一次 c：
| β | c 精确 | JAX jvp(grad) | tenferro 差分 | 两边之差 | tenferro − 精确 |
|---|---|---|---|---|---|
| 0.4 | 0.861698 | 0.861447344 | 0.861447346 | +2.3e-09 | -0.0003 |
| 0.42 | 1.197468 | 1.196542125 | 1.196542118 | -6.1e-09 | -0.0009 |
| 0.43 | 1.526848 | 1.524985776 | 1.524985783 | +7.0e-09 | -0.0019 |
| 0.435 | 1.841033 | 1.839392416 | 1.839392424 | +8.1e-09 | -0.0016 |
| 0.44 | 2.889687 | 3.311956763 | 3.311956765 | +2.8e-09 | +0.4223 |
| 0.445 | 1.985281 | 1.981538637 | 1.981538599 | -3.8e-08 | -0.0037 |
| 0.45 | 1.607953 | 1.607084462 | 1.607084499 | +3.7e-08 | -0.0009 |
| 0.46 | 1.253571 | 1.253240620 | 1.253240623 | +2.9e-09 | -0.0003 |
| 0.48 | 0.914803 | 0.914598430 | 0.914598431 | +1.5e-09 | -0.0002 |
两边最多差 4e-8：导数末位约 1e-13 的差，除以 2h = 2e-6 再乘 β²，就是 1e-8 量级。这比 χ = 16 的物理误差小三个量级以上。在所取的 β 里，c 的峰值落在 β = 0.44，比精确解高 0.42；精确解在 β_c 处发散，有限 χ 的 TRG 把它抹成一个有限高的峰。
❗IMPORTANT
动脑筋：把 χ 从 16 加到 24，β = 0.44 的内能 u 误差（−5.58e-4）会缩小几倍？
只改命令行的第一个参数：ising_trg 24 20 8 1e-6 0.44。截断数变成 [2, 4, 4, 16, 24, 24, …]，最大的矩阵从 256×256 变成 576×576。先猜，再看文末。
换成 traced 分段
同样的分段也可以用 traced 写：每种步的形状 (d, k) 建一个小程序，β、T 和余切都是占位符，编译、准备一次，20 步、所有 β 复用。实验代码是 bin trg_staged，建程序的那一段：
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;
}
bwd 是对标量替身 ⟨T′, ct⟩ + c·ln‖T‖ 求 grad，一个程序同时吃 T、ct 和 c。结果和 eager 一致（β = 0.30 的 u 是 −0.704498679529754，eager 是 −0.704498679529762），时间不同。每个 β 的「求值加梯度」，稳态：
| 程序 | 1 线程 | 8 线程 |
|---|---|---|
JAX jit(value_and_grad)（XLA 默认线程设置，列在右栏） | — | 188–226 ms |
tenferro eager 分段，逐步 vjp（本页主程序） | 540–558 ms | 266–275 ms |
tenferro eager 分段，标量替身加 backward | 541–560 ms | — |
| tenferro eager 分段，收缩写成两两相乘的链 | 630–650 ms | 308–380 ms |
| tenferro traced 分段 | 887–923 ms | 440–469 ms |
（JAX 是第一次调用之后 5 次的最小值，第一次调用含追踪和编译 1.23 s；tenferro eager 是第一次之后 3 次的最小值；traced 是 5 次的最小值，所有程序的编译加准备共 27.5 ms。五个 β 给出范围。）
主程序选 eager，有三个理由，都是这次量出来的或写在代码里的：
- 快。 1 线程 eager 比 traced 少用约 40% 的时间。traced 的每个
bwd程序 204 个运算，前向程序只有 40 个：反向要把这一步的前向、包括两次 SVD 重算一遍；eager 的vjp直接用前向留下的残差（eager 还是 traced）。 - 短，而且和 JAX 一一对应。
step就是 JAX 循环体的逐行翻译，contraction 能用 einsum；traced 要另写mk把每个输出编译、准备成程序，按 (d, k) 存进表，run时按顺序传输入。 vjp和backward一样快。 逐步vjp和标量替身加backward两种写法时间相同，导数差几个 ulp（各自加法的顺序不同）；vjp不碰梯度槽，读起来更接近jax.vjp。
traced 分段的长处在别处：所有程序在开头一次建好，之后每次只是 run，没有 eager 每次反向的簿记；程序和形状绑定，形状不对立即报错。怎么选见 分段；时间怎么量见 计时。
和 JAX 比，tenferro 8 线程的 eager 分段慢 1.2–1.5 倍，1 线程慢 2.4–2.9 倍。收缩的写法也有影响：写成两两相乘的链，1 线程慢约 15%；四操作数的 einsum 让规划器选 (S1·S2)(S3·S4)，和 JAX 的路径相同。一步里大头是两次 256×256 的 SVD，所以收缩树的差别只有这么多。JAX 的整张图交给 XLA 编译，20 步之间没有 Python 或宿主的往返；tenferro 每步都要把张量读回宿主再做成新叶子。
动脑筋答案
缩小约 8.5 倍：
tenferro 0.8.0 eager staged chi=24 steps=20 width=8 h=1e-6 ks=[2, 4, 4, 16, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24]
beta=0.4400 lnZ/N=0.928727747415729 u=-1.402292812433066 c=3.036963245040 | value_and_grad 1446.9 ms
- u 的误差从 −5.583e-4 变成 −1.402292812433066 − (−1.402226960031489) = −6.585e-5，小了 8.5 倍；lnZ/N 的误差从 −4.578e-6 变成 −7.930e-7，小了 5.8 倍。
- c 从 3.3120 降到 3.0370，误差从 +0.42 降到 +0.15：峰还在，变矮了。χ 越大，有限 χ 抹平的发散恢复得越多。
- 一次求值加梯度从约 0.27 s 涨到 1.45 s（8 线程）：最大的 SVD 从 256×256 变成 576×576。JAX 同样的 χ = 24：lnZ/N 差 5.6e-16，u 差 1.4e-14，
jit(value_and_grad)稳态 1.06 s。
小结
- 同一个 TRG：JAX 一个函数、Python 循环展开 20 步，
jit(value_and_grad)；tenferro 每步一张 eager 图，variable_from切断 trace，反向逐步两次rt.vjp串起余切。step是 JAX 循环体的逐行翻译，einsum（包括取迹的xyxy->）下标串相同。 - JAX 一侧要自己用
custom_jvp写展宽的 SVD 规则，tenferro 的默认规则就是它；reshape 用order="F"让矩阵化一致；截断数前三步压到结构性的秩 2、4、4。 - 两边 lnZ/N 差 7.8e-16、导数差 5.2e-13（tenferro 1 线程）；和 Onsager 的差是 χ = 16 的截断误差，两边相同。
- 比热用 AD 能量的中心差分。β = 0.44 时 h = 1e-5 跨过了最后一步截断处第 16、17 个奇异值的交叉，lnZ/N 比 AD 斜率多涨 3.7e-9，c 偏 1.3%；h = 1e-6 时和 JAX 的
jvp(grad)一致。用两个 h 交叉检查。 - 时间：tenferro eager 分段 1 线程约 0.55 s、8 线程约 0.27 s 一个 β，比 traced 分段快（反向用残差，不重算 SVD），比 JAX 的
jit（约 0.2 s）慢。
💬 关联讨论
发起讨论加载中...