einsum

TRG 每一步把四个三阶张量收成下一步的张量，JAX 里是一行 jnp.einsum。这一部分只有一页，拿这一步的收缩对照 jnp.einsum 和 tenferro：tenferro 按模式分了三个 einsum 入口，TracedTensor 上没有 einsum，traced 的 TRG 要自己写 dot_general。einsum 的解析、规划、执行和求导规则见 tenferro 卷的 einsum。

1. einsum：两边对 war,abu,bgl,gwd->ruld 选同一棵树，tenferro 的三个入口和 JAX 逐位相同；optimize=False 两边同义，JAX 的路径列表可以原样交给 EinsumOptimize::Path；手写 dot_general 时照 einsum_path 的树比链快一倍；dot_general 的批维在最后；TraceContext 的 einsum 用 AdContext::vjp_program 求导；长度 1 广播时 eager 的梯度形状错。

flowchart TB
  j["jnp.einsum"] --> c["具体张量 einsum"]
  j --> e["eager einsum"]
  j --> tc["TraceContext einsum"]
  j --> dg["TracedTensor dot_general"]
  tc --> vp["vjp_program 求导"]
  e --> bw["backward 求导"]
  dg --> gr["ad.grad 求导"]