Rust in Everything
基于 tenferro-einsum 0.8.0

einsum

主线位置：TRG 主线「一步 SVD 截断加收缩」里的收缩：四个三阶张量收成下一步的 T，war,abu,bgl,gwd->ruld。

版本：JAX 0.12.0.1（CPU，x64）、opt_einsum 3.4.0、tenferro-rs 0.8.0（CPU，native）
flowchart TB
  j["jnp.einsum"] --> p["opt_einsum 定路径"]
  p --> jx["jaxpr 里的 dot_general"]
  c["具体张量 einsum"] --> pl["tenferro 定树"]
  e["eager einsum"] --> pl
  tc["TraceContext einsum"] --> ext["一个扩展算子"]
  ext -->|"运行时"| pl
  tt["TracedTensor"] --> dg["dot_general 自己选树"]

JAX 只有一个 jnp.einsum：在 jit 里、grad 里、普通调用里都是它。追踪时 opt_einsum 选好两两收缩的路径，jaxpr 里留下的是一串 dot_general 加一次 transpose，之后的变换只看见这些原语。tenferro 按模式分了入口：具体张量（[&a, &b].einsum(..)）、eager 会话里的 s.einsum(..)、traced 的 TraceContext::einsum。TracedTensor 上没有 einsum，traced 程序里最常见的写法是直接写 dot_general。规划、执行和求导规则的内部见 tenferro 卷的 einsum，这一页只对照 JAX 的用法。

实验代码：tenferro 一侧是 bin einsum_tf（release），JAX 一侧是 einsum_jax.py。M2 Max、macOS 27、rustc 1.99.0；uptime 的 1 分钟负载 4.8–5.7（同时有别的实验在跑）。四个输入都是 D×D×D，S_n[a,b,c] = sin(1 + n + 0.37·(a + D·b + D²·c))，括号里是列主序的平铺下标，所以两边的数据完全一样（列主序见 布局与 reshape）。除了计时一节，D = 4，CpuEngine::serial()。

同一个收缩，三个入口

t = jnp.einsum("war,abu,bgl,gwd->ruld", s1, s2, s3, s4)

const SPEC: &str = "war,abu,bgl,gwd->ruld";
// 具体张量：直接算，不记录
let concrete = [&st[0], &st[1], &st[2], &st[3]].einsum(SPEC, &mut eng)?;
// eager：在会话里立即执行并记录
let et = ert.with_eager_session(&mut eng, |ss| -> R<EagerTensor> {
  Ok(ss.einsum(&[&es[0], &es[1], &es[2], &es[3]], SPEC)?)
})?;
// traced：TraceContext 里的值是 TraceValue，不是 TracedTensor
let mut tc = TraceContext::new();
let mut vals = Vec::new();
for n in 0..4 {
  vals.push(tc.input_with_default(ProgramInputSpec::new(DType::F64, DimExpr::from_concrete(&[d, d, d])), s_tensor(n, d)?)?);
}
let t = tc.einsum_with(&vals, SPEC, EinsumOptimize::default())?;
let prog = GraphCompiler::new().compile_traced_graph(&tc.finish(&[t])?)?;

# JAX
shape (4, 4, 4, 4) sum -57.97394471558693 T[0,1,2,3] 0.35516360216185444
# tenferro
concrete einsum: shape [4, 4, 4, 4] max|. - jax| = 0.000e0 bitwise 256/256
eager einsum: max|. - jax| = 0.000e0 bitwise vs concrete 256/256
TraceContext einsum (default): max|. - jax| = 0.000e0 bitwise vs concrete 256/256
TraceContext einsum (Path from jnp.einsum_path): max|. - jax| = 0.000e0 bitwise vs concrete 256/256
TraceContext einsum (False): max|. - jax| = 0.000e0 bitwise vs concrete 256/256
max|T| = 0.8472
  • 结果。 tenferro 的三个入口、三种计划，256 个元素和 JAX 全部逐位相同（|T| 最大 0.85）。下一节会看到两边选的是同一棵收缩树，这是逐位相同的前提。
  • 下标要写箭头。 字母当标签两边一样；jnp.einsum("ij,jk", a, b) 按 NumPy 规则推出输出 (2, 4)，tenferro 不写箭头时报 einsum-arrow。
  • traced 入口不在 TracedTensor 上。 TraceContextEinsumExt 是 TraceContext 的方法，收发的是 TraceValue；TracedTensor 上的 TracedTensorEinsumExt 只有 tensordot。两种句柄不能混用，后面两节分别讲这带来的写法和求导问题。
  • 运行时需要 einsum 模块。 traced 的 einsum 在程序里是一个扩展算子，Runtime 要装 tenferro_einsum::extension_module，不装的报错见 tenferro 卷 einsum 的 traced 一节。

收缩路径：einsum_path 和 ContractionTree

for opt in ["auto", "greedy", "optimal", False]:
  path, info = jnp.einsum_path(SPEC, *s, optimize=opt)

let subs = Subscripts::parse(SPEC)?;
let tree = ContractionTree::optimize(&subs, &[&sh[..], &sh[..], &sh[..], &sh[..]])?;
for i in 0..tree.step_count() {
  let (l, r, o) = tree.step_subscripts(i).ok_or("step")?;
  println!("tenferro plan step {i}: pair {:?}  {},{}->{}", tree.step_pair(i).ok_or("pair")?, labels(l), labels(r), labels(o));
}

# JAX
optimize='auto': path=[(0, 1), (0, 1), (0, 1)] flops=12288 largest=256
optimize='greedy': path=[(0, 1), (0, 1), (0, 1)] flops=12288 largest=256
optimize='optimal': path=[(0, 1), (0, 1), (0, 1)] flops=12288 largest=256
optimize=False: path=[(0, 1), (0, 1), (0, 1)] flops=12288 largest=256
...
   5           TDOT          abu,war->buwr                    bgl,gwd,buwr->ruld
   5           TDOT          gwd,bgl->wdbl                       buwr,wdbl->ruld
   6           TDOT        wdbl,buwr->ruld                            ruld->ruld
{ lambda ; a:f64[4,4,4] b:f64[4,4,4] c:f64[4,4,4] d:f64[4,4,4]. let
    e:f64[4,4,4,4] = dot_general[
      dimension_numbers=(([0], [1]), ([], []))
...
    h:f64[4,4,4,4] = transpose[permutation=(3, 2, 1, 0)] g
  in (h,) }
# tenferro
tenferro plan step 0: pair (0, 1)  war,abu->wrbu
tenferro plan step 1: pair (2, 3)  bgl,gwd->blwd
tenferro plan step 2: pair (4, 5)  wrbu,blwd->ruld
TraceContext program: inputs=4 ops=1
  SemanticOpRef::Extension("tenferro.einsum.v1")

两种写法记的是同一棵树。opt_einsum 的路径是「位置」：从操作数列表里取出第 i、j 个，结果追加到末尾，所以 [(0, 1), (0, 1), (0, 1)] 是先收 S1·S2，再收 S3·S4，最后两个中间结果相收。tenferro 的 step_pair 用的是编号：输入 0–3，中间结果从 4 往后编，(0, 1)、(2, 3)、(4, 5) 是同一个顺序。TRG 这一步四种 optimize 的结果一样，(S1·S2)(S3·S4) 本来就是最省的树。

两边的差别在「什么时候定」。JAX 在追踪时定路径，jaxpr 里已经是三个 dot_general 和一个 transpose。tenferro 的 traced 程序里只有一个 Extension("tenferro.einsum.v1")，树等到运行时知道形状才建，按形状缓存（einsum 的 traced 一节）。

optimize 参数的对应：

JAX optimize=tenferro EinsumOptimize::说明
"auto"（默认）、"greedy"、"optimal"Auto(..)（默认）tenferro 默认是自己的确定性贪心；开了退火才用 omeco 的 TreeSA，见 einsum 的规划一节
FalseFalse两边都是每步取位置 (0, 1)，四个以上输入时是两两配对，不是从左到右
路径列表 [(0, 1), ...]Path(vec![(0, 1), ...])jnp.einsum_path 的结果原样搬过来
没有Nested（"(ij,jk),kl->il"）、Tree括号指定顺序、预先算好的树

只有 traced 的 einsum_with 收 EinsumOptimize；具体张量和 eager 的 einsum 没有这个参数，只用默认计划。Path 的文档写的就是「JAX-compatible position-based contraction path」。

False 这一行容易让人以为是从左到右。jax 0.12.0.1 的 jax/_src/numpy/einsum.py 里 optimize=False 用的是：

class Unoptimized(opt_einsum.paths.PathOptimizer):
  """Unoptimized path for einsum."""
  def __call__(self, inputs, *args, **kwargs):
    return [(0, 1)] * (len(inputs) - 1)

每步取当前列表的前两个、结果放到末尾，和 tenferro 的 False 是同一个规则。所以从 JAX 搬过来的 optimize=False 在 tenferro 里语义不变；它在长链上会造出外积，见常见坑 einsum-optimize-false。

traced 张量上：dot_general

在 traced 模式里写 TRG，局部张量、SVD、切片都是 TracedTensor 的方法，收缩也得在 TracedTensor 上做，TraceContext 的 einsum 接不进来。办法是自己写 dot_general（或 tensordot），把 einsum 的一步翻成收缩轴：

x1 = jnp.einsum("war,abu->wrbu", s1, s2)
x2 = jnp.einsum("wrbu,bgl->wrugl", x1, s3)
t = jnp.einsum("wrugl,gwd->ruld", x2, s4)

// 一条链：((S1 S2) S3) S4，和上面三行 JAX 一样
let x1 = s[0].dot_general(&s[1], dg(&[1], &[0], &[], &[]))?;     // war,abu->wrbu
let x2 = x1.dot_general(&s[2], dg(&[2], &[0], &[], &[]))?;       // wrbu,bgl->wrugl
let t = x2.dot_general(&s[3], dg(&[3, 0], &[0, 1], &[], &[]))?;  // wrugl,gwd->ruld

// 照 einsum_path 的树：(S1 S2)(S3 S4)
let x12 = s[0].dot_general(&s[1], dg(&[1], &[0], &[], &[]))?;    // war,abu->wrbu
let x34 = s[2].dot_general(&s[3], dg(&[1], &[0], &[], &[]))?;    // bgl,gwd->blwd
let t = x12.dot_general(&x34, dg(&[0, 2], &[2, 0], &[], &[]))?;  // wrbu,blwd->ruld

dg(lhs 收缩轴, rhs 收缩轴, lhs 批轴, rhs 批轴) 是填 DotGeneralConfig 的小函数。没有批轴时，dot_general 的输出是「左边剩下的轴，再接右边剩下的轴」，各自保持原顺序，和 lax.dot_general 一样，所以每行注释里的输出下标可以直接读出来。tensordot 是同一件事的简写：tensordot 把轴对翻成一个没有批轴的 DotGeneralConfig，再调 dot_general。

第一种写法是试验阶段 TRG 用的链。手写的时候，树由你决定，没有人替你选。

❗IMPORTANT
动脑筋：第一种写法（链）的 256 个结果里，有几个和 jnp.einsum 四操作数的结果逐位相同？第二种呢？两种写法在数学上相等，单线程。

dot_general 的输出轴序：批维在最后

有批轴时两边不一样：

lhs = jnp.arange(24.0).reshape(2, 3, 4); rhs = jnp.arange(40.0).reshape(2, 4, 5)
jax.lax.dot_general(lhs, rhs, (((2,), (1,)), ((0,), (0,))))

let o = lt.dot_general(&rt2, dg(&[2], &[1], &[0], &[0]))?;   // 同样的数据，列主序装入

# JAX
lhs (2,3,4) rhs (2,4,5), contract 2/1, batch 0/0 -> shape (2, 3, 5)
lhs (3,4,2) rhs (4,5,2), contract 1/0, batch 2/2 -> shape (2, 3, 5)
# tenferro
lhs [2,3,4] rhs [2,4,5], contract 2/1, batch 0/0 -> shape [3, 5, 2]
out[i,j,b] == jax out[b,i,j]: max diff 0

lax.dot_general 的输出总是「批、左自由、右自由」，批轴在输入里放在哪儿都一样。tenferro 是「左自由、右自由、批」，批轴在最后，数值完全相同（差 0），只是轴的位置不同。这和 tenferro 批量线代的批维在右边是同一个约定，见 布局与 reshape 的批维一节。把 jaxpr 里的 dot_general 照搬过来时，后面的 transpose 要按这个顺序重写。einsum 不受影响，输出下标写什么就是什么。

求导

JAX 里 einsum 只是一串 dot_general，jax.grad 照常走。tenferro 有三条路，结果都对，前提各不相同。损失是 L = Σ T²，对 S1 求导：

g = jax.grad(lambda s1, s2, s3, s4: jnp.sum(jnp.einsum(SPEC, s1, s2, s3, s4) ** 2))(*s)

// 1. traced：dot_general 链加 ad.grad
let l = (&t * &t)?.reduce_sum(None)?;
let g = ad.grad(&l, &s[0])?;

// 2. eager：einsum 加 backward，EagerRuntime::new() 不带任何规则
let l = rt.with_eager_session(eng, |ss| -> R<EagerTensor> {
  let t = ss.einsum(&[&s[0], &s[1], &s[2], &s[3]], SPEC)?;
  let tt = ss.mul(&t, &t)?;
  Ok(ss.reduce_sum(&tt, None)?)
})?;
l.backward(&mut eng)?;

// 3. TraceContext：没有 TracedTensor，走程序级的变换
let tt = tc.add_op(CoreSemanticOp::Mul, &[t, t])?[0];
let l = tc.add_op(CoreSemanticOp::ReduceSum { axes: vec![0, 1, 2, 3] }, &[tt])?[0];
let prog = GraphCompiler::new().compile_traced_graph(&tc.finish(&[l])?)?;
let sem = ad.vjp_program(prog.frozen_program(), &[true, false, false, false], &[true])?;
let gprog = GraphCompiler::new().compile_frozen_program(sem.frozen())?;
// run 时输入是 [S1, S2, S3, S4, 种子 1.0]

# JAX
L = 43.28075539488678 shape (4, 4, 4) sum 42.33930479738713 max|G| = 35.03450974377743
G[1,2,3] = 18.330572446790242 FD = 18.330572455482752
# tenferro
traced dot_general x3 + ad.grad: ops=12 max|. - jax| = 1.066e-14  G[1,2,3] = 18.33057244679025
eager einsum + backward, EagerRuntime::new(): L = 43.28075539488680 max|. - jax| = 7.105e-15 bitwise vs traced 9/64
eager einsum + backward, with einsum rules: L = 43.28075539488680 max|. - jax| = 7.105e-15 bitwise vs traced 9/64
TraceContext loss program: inputs=4 outputs=1
AdContext without rules: vjp_program Err(semantic extension AD failed: semantic extension family "tenferro.einsum.v1" has no Linearize AD rule)
AdContext with einsum rules: vjp_program ok, inputs=5 outputs=1 seed index [Some(4)] grad index [Some(0), None, None, None]
  SemanticOpRef::Extension("tenferro.einsum.v1")
  SemanticOpRef::Core(<bounded>)
  SemanticOpRef::Core(<bounded>)
  SemanticOpRef::Core(<bounded>)
  SemanticOpRef::Core(<bounded>)
  SemanticOpRef::Extension("tenferro.einsum.v1")
  grad: max|. - jax| = 1.243e-14 bitwise vs traced dot_general 7/64
  prepare with no inputs: Err(placeholder semantic input 4 has no runtime input binding)
  • 数值。 三条路和 JAX 的梯度差 7e-15 到 1.2e-14，梯度最大 35，相对 3e-16 量级；JAX 自己对中心差分的一个元素差 8.7e-9（差分步长 1e-6 的截断误差）。三条路彼此也不逐位相同（64 个里 7 到 9 个），各自的反向计算顺序不同。
  • traced 链：全是核心算子，不需要 einsum 的规则。导数程序 12 个算子。
  • eager：不带规则的 EagerRuntime::new() 和带规则的结果相同。
  • TraceContext：ad.grad 只收 TracedTensor，用不上；AdContext::vjp_program 直接变换整个程序，从 CompiledGraph::frozen_program 拿输入，结果用 compile_frozen_program 编译。两个掩码选「对哪些输入求导」「哪些输出给种子」。种子追加在原输入之后（第 4 个），没有默认值，run 时必须给，和 grad、vjp、jvp 里 ad.grad 的种子默认是 1 不同。导数程序里有两个 einsum 扩展算子，反向收缩也是 einsum，运行时同样要装模块。这条路要 AdContext 里有 einsum 的规则，没有就报 has no Linearize AD rule。
📝NOTE
没有蠢问题

问： 同一个 einsum，eager 不带规则能 backward，TraceContext 的程序为什么非要 einsum 的规则？

答： 因为 eager 记下的根本不是 einsum。einsum_subscripts_with_broadcast 在不需要长度 1 广播时走 try_expand_eager_einsum，把收缩展开成核心算子逐个执行、逐个记录，反向用核心算子的规则。traced 的程序里留着一个完整的 einsum 扩展算子，变换走到它只能查扩展规则。eager 只有在需要长度 1 广播时才记一个 einsum 扩展算子，那时才用到 einsum 的规则，正好碰上下一节的问题。

长度 1 的广播

NumPy 和 JAX 的 einsum 允许同一个标签一边长 1、一边长 n，按广播算：

a = np.array([[1.0], [2.0]]); b = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
jnp.einsum("ij,jk->ik", a, b)
jax.grad(lambda a: jnp.sum(jnp.einsum("ij,jk->ik", a, b)))(a)

# JAX / NumPy
numpy einsum: [[9.0, 12.0], [18.0, 24.0]]
jnp einsum: [[9.0, 12.0], [18.0, 24.0]]
jax grad shape (2, 1) [[21.0], [21.0]]
# tenferro
concrete: shape [2, 2] [9.0, 18.0, 12.0, 24.0]
eager: sum = [63.0], grad shape [2, 3] [3.0, 3.0, 7.0, 7.0, 11.0, 11.0]
traced: Err(extension family "tenferro.einsum.v1" shape constraint at instruction None failed: Const(1) (1) Equal Const(3) (3))

前向：具体张量和 eager 的结果和 JAX 相同（列主序打印），traced 不写 ... 时在编译时拒绝。反向：eager 给出 [2, 3] 的梯度，形状是广播之后的，JAX 是 (2, 1)、[[21], [21]]。这是 0.8.0 einsum 反向规则的已知问题，没有沿广播轴求和，见常见坑 einsum-broadcast-grad 和 einsum-broadcast-forward。从 JAX 搬代码时，靠长度 1 广播的 einsum 要先 broadcast_in_dim 成同样的长度再收缩。

时间

D = 16，四个 16×16×16 收成 16⁴，最省的树 3.8e7 次浮点运算。JAX 一侧 200 次、tenferro 一侧 50 次，取最小值和中位数（ms）；JAX 的首次调用含追踪和编译，tenferro 的 traced 程序建图、编译、准备在 0.1–0.5 ms 之间，不算在下表里。

JAX jittenferro 单线程tenferro 8 线程
四操作数 einsum，JAX 首次 27 ms0.362 / 0.378具体 1.610 / 1.737；TraceContext 1.534 / 1.594；eager 1.516 / 1.533具体 0.703 / 0.941；TraceContext 0.558 / 0.672；eager 0.478 / 0.644
链：三次二元收缩0.571 / 0.625dot_general 3.270 / 3.6370.972 / 1.258
dot_general 照 einsum_path 的树—1.511 / 1.5700.484 / 0.630
Σ T² 对 S1 的梯度0.648 / 0.700链 6.491 / 7.030；照树 2.993 / 3.243链 2.419 / 3.335；照树 1.176 / 1.392
  • 链比树慢一倍。 链的第二步 wrbu,bgl->wrugl 造出一个 16⁵ 的中间结果，三步一共约 6.9e7 次浮点运算，是最省的树的 1.8 倍；tenferro 单线程的时间比是 3.27 / 1.51 = 2.2，JAX 是 0.571 / 0.362 = 1.6。手写 dot_general 时照 jnp.einsum_path 的路径写，前向和梯度都快一倍。试验阶段 TRG 的每一步用的是链。
  • JAX 不是单线程。 2000 次调用的进程 CPU 时间是墙钟的 2.6–2.7 倍（测了两轮）；XLA_FLAGS="--xla_cpu_multi_thread_eigen=false intra_op_parallelism_threads=1" 不改变这个比例，也不改变时间。所以更可比的是 tenferro 8 线程那一列：同一棵树，tenferro 0.48–0.70 ms，JAX 0.36 ms。
  • 三个 einsum 入口时间相近。 TraceContext 的第一次 run 要规划树，之后按形状查缓存。

动脑筋答案

traced dot_general x3: max|. - jax| = 9.905e-16 bitwise vs concrete 16/256
traced dot_general pairwise (S1 S2)(S3 S4): max|. - jax| = 0.000e0 bitwise vs concrete 256/256
# JAX 自己的两种写法
max|4-operand - three binary| = 9.905271047827569e-16 max|T| = 0.8472346948742482
  • 链：16 / 256 逐位相同，最大差 9.9e-16（|T| 最大 0.85）。树不同，每个元素的乘加顺序跟着不同，舍入就不同。
  • 照 einsum_path 的树：256 / 256，和 tenferro 的 einsum、JAX 的四操作数 einsum 完全一样。
  • JAX 自己的链和四操作数写法也差 9.9e-16，和 tenferro 的链对四操作数的差一样：同一棵树给同样的位，换树才变，和用哪个库无关。

小结

  • jnp.einsum 一个函数到处能用；tenferro 分具体张量、eager、TraceContext 三个入口，TracedTensor 上只有 tensordot / dot_general。
  • TRG 这一步两边选同一棵树 (S1·S2)(S3·S4)，tenferro 三个入口和 JAX 的结果 256 / 256 逐位相同。
  • optimize 的对应：False 两边同义（每步取位置 (0, 1)），JAX 的路径列表原样交给 EinsumOptimize::Path；只有 traced 的 einsum_with 收这个参数。
  • 在 TracedTensor 上手写 dot_general 时树由自己定：照 einsum_path 的树比链快一倍，结果和 einsum 逐位相同。
  • dot_general 的输出是「左自由、右自由、批」，lax.dot_general 是「批、左自由、右自由」。
  • 求导：dot_general 链用 ad.grad；eager einsum 展开成核心算子，不要规则；TraceContext 的 einsum 用 AdContext::vjp_program，要 einsum 规则，种子是没有默认值的最后一个输入。
  • 长度 1 广播：前向具体张量和 eager 能算、traced 拒绝，eager 的梯度形状错（0.8.0 的已知问题），先显式广播。

💬 关联讨论

发起讨论
加载中...