compute_eta_jets#
- t3toolbox.backend.sampling_derivatives.compute_eta_jets(tt_cores, mu_jets, nu_jets, trs)#
def compute_eta_jets( tt_cores: typ.Sequence[NDArray], # len=d, elm_shape=C+(rLi,nOi,rR(i+1)) mu_jets: typ.Sequence[NDArray], # len=d, elm_shape=(order+1,)+W+C+(rLi,) nu_jets: typ.Sequence[NDArray], # len=d, elm_shape=(order+1,)+W+C+(rR(i+1),) trs: NDArray, # binomial tensor, shape=(order+1,order+1,order+1) ) -> typ.Tuple[NDArray, ...]: # eta_jets. len=d, elm_shape=(order+1,)+W+C+(nOi,)
Combine the left and right jets at each free mode (standard order-scan form).
Dense reference:
compute_eta_jets_trs()(equal to tolerance).eta is a FULL convolution (both jets run
0..order), so the densetrscontraction forms the spatial productmu . Gwith mode AND bond live at every order at once: an(order+1)*W*r^2intermediate (and the uniform d-einsum forms it for every core at once –dtimes worse, 128 GB at r=128, W=32000, order 5). This scans the two small axes instead, so only one slice is ever live:eta^(t) = sum_r C(t,r) [ mu^(r) . G ] . nu^(t-r)xscanover the input orderr(the convolution): formsmu^(r) . Gone order at a time and accumulates – removes the(order+1)factor;xmapover cores: removes thedfactor.
The win requires the UNIFORM path (a stacked supercore), where
xmap/xscandispatch to the realjax.lax.map/jax.lax.scan– both sequential (lax.mapis ascan, not a vectorizingvmap), so only one core+order intermediate is resident. On the RAGGED pathxmap/xscanare Python loops that unroll under jit and keep everything co-resident (measured ~1.2x, i.e. no win). Uniform, both scans: measured 14-28x smaller XLA peak than the dense d-einsum and CONSTANT in order (~4.5 GB vs 64-128 GB at the huge config). The order loop being a real scan is what forces this – an unrolled loop does not. Verified equal tocompute_eta_jets_trs()to 1e-12 (ragged) and 1e-7 (uniform, float32).