Module recalculated 

Physical-frame moments with a custom, memory-efficient backward

The forward is the same serial chunkwise computation as Mamba3MomentsInput::state_moments_phys, but routed through the Mamba3MomentsBackendExt trait so Autodiff backends substitute a custom recompute backward (see super::backward): the node saves only its five leaf inputs and the tiny (m2, m1) outputs; backprop re-materialises one chunk of states at a time — the full per-token state trajectory that plain autodiff would retain is never kept alive. This is the at-scale execution model.

The default body (every plain backend) runs the identical math on B’s primitives via the rank-tagged F wrapper. The rotation is passed as a rank-erased primitive plus a quaternion flag ([b, s, h, a] cumulative angles, or [b, s, h, J, 4] cumulative unit quaternions), so a single trait method covers both algebras.

Enums

RotPrim 🔒

The per-token cumulative rotation on B’s primitives — the rank-tagged counterpart of RotationSeq.

Traits

Mamba3MomentsBackendExt

Extends the backend with the physical-frame state-moments computation.

Functions

chunk_cos_sin 🔒

The rotation’s per-pair cos/sin for one chunk, expanded over per_head_dim: [b, h, l, p, a] each. transpose negates the angle — false applies the de-rotation R(−θ) (cache → physical), true its transpose R(+θ) (used on cotangents).

chunk_decay 🔒

Chunk-local decay: dₜ = exp(a_cum[t]) ([b, h, l]) and the 1-semiseparable mask L[t, j] = exp(a_cum[t] − a_cum[j]) ([b, h, l, l], zero strictly above the diagonal).

chunk_mask 🔒

Float validity mask [1, 1, l, 1, 1] for a chunk starting at start (1 for global positions < valid_len, 0 for pads).

chunk_states 🔒

One chunk of cache-frame states h̃[t] = dₜ·h₋ + Σ_{j≤t,c} L[t,j]·x̂[j,c] ⊗ b̂[j,c] — the channel axis is folded into the write index for one batched GEMM.

chunk_writes 🔒

The per-token write W_j = Σ_c x̂[j,c] ⊗ b̂[j,c] ([b, h, j, p, r]) — used by the backward’s da gradient.

moments_phys_fwd 🔒

The serial chunkwise forward on primitives — the exact math of Mamba3MomentsInput::state_moments_phys.

quat_conj_prim 🔒

Quaternion conjugate on the packed trailing axis.

quat_mul_prim 🔒

Quaternion Hamilton product on the packed trailing (w, x, y, z) axis (primitive port of rotation::quat_mul).

rotate_chunk 🔒

Rotate the state_rank axis of one chunk of states by the per-token cumulative rotation. transpose = false is the de-rotation to the physical frame (R(−θ) / left-multiply by Q); transpose = true is its transpose (R(+θ) / left-multiply by conj(Q)), the map cotangents take back to the cache frame.