Module backward 

Custom autodiff node for the physical-frame state moments

Implements Mamba3MomentsBackendExt for Autodiff<B> via a single Burn [Backward] node. The forward saves only its five leaf inputs; backprop re-materialises one chunk of states at a time (the SerialRecalculated discipline) and evaluates the analytic VJPs, so the full per-token state trajectory plain autodiff would retain is never kept alive.

Gradient math (per chunk, reverse order, cotangent carry d_h₊)

Forward: h̃[t] = dₜ·h₋ + Σⱼ L[t,j]·Wⱼ with writes Wⱼ = Σ_c x̂ⱼ_c ⊗ b̂ⱼ_c, cₜ = Rₜ†·h̃[t] (de-rotation), m2 += Σₜ maskₜ·cₜᵀcₜ, m1 += Σₜ maskₜ·cₜ.

• d_c[t] = maskₜ·(cₜ·(d_m2 + d_m2ᵀ) + d_m1) — the moments’ VJP.
• d_h̃[t] = Rₜ·d_c[t] (the de-rotation’s transpose), plus the rotation’s own gradient: per complex pair, ∂cₜ/∂θ = (c_y, −c_x), so d_θ = Σ_p (d_c_x·c_y − d_c_y·c_x); per quaternion block (c = Q ⊗ h̃), d_Q = Σ_p d_c ⊗ conj(h̃) (bilinear-product VJP, exact for all quaternions). This is what lets a PR penalty shape the rotation itself.
• G[t] = d_h̃[t] + δ_{t,l−1}·d_h₊ (the next chunk consumed h̃[l−1] as its h₋).
• d_h₋ = Σₜ dₜ·G[t] (the new carry), d_Wⱼ = Σₜ L[t,j]·G[t], then d_x̂ⱼ_c = d_Wⱼ·b̂ⱼ_c, d_b̂ⱼ_c = d_Wⱼᵀ·x̂ⱼ_c.
• d_da: through dₜ = exp(Σ_{s≤t} da_s) and L[t,j] = exp(Σ_{j<s≤t} da_s), with scalₜ = ⟨G[t], h₋⟩·dₜ and A[t,j] = ⟨G[t], Wⱼ⟩·L[t,j]: d_da_s = Σ_{t≥s} scalₜ + Σ_{j<s≤t} A[t,j] = rev_cumsum(scal + Σⱼ A[·,j])_s − rev_cumsum(Σₜ A[t,·])_s (the pair (t, j) contributes to exactly the positions j < s ≤ t, and 1[s≤t] − 1[s≤j] is that indicator; A is lower-triangular).
• The boundary carries h₋(n) are recomputed by a cheap chunk-level pass (decay-to-chunk-end GEMM, l× cheaper than the states GEMM).

Chunks entirely past valid_len were never entered by the forward (break), so their gradients are zero and the reverse loop skips them.

The two outputs (m2, m1) are flattened into one tracked tensor via crate::utils::combined_grad so a single Backward<B, 5> node covers both.

Structs

MomentsGrads 🔒

All input gradients of the moments node.

Functions

boundary_carries 🔒

Recompute the boundary carries h₋(n) entering each of the processed chunks (a chunk-level pass: decay-to-chunk-end GEMM + serial accumulation).

d_quat_chunk 🔒

The per-block quaternion gradient of the de-rotation (c = Q ⊗ h̃ ⇒ d_Q = Σ_p d_c ⊗ conj(h̃)).

d_theta_chunk 🔒

The per-pair angle gradient of the de-rotation (see the module header).

moments_phys_bwd 🔒

The full recompute backward (see the module header).

rev_cumsum 🔒

Suffix sum along dim (out[s] = Σ_{t≥s} in[t]).

scatter_last 🔒

Cotangent scatter: [b,h,p,r] placed at the last position of a [b,h,l,p,r] chunk (zeros elsewhere).