Struct Mamba3MomentsInput 

pub struct Mamba3MomentsInput {
    pub xhat_bnlMhp: Tensor<6>,
    pub bhat_bnlMhr: Tensor<6>,
    pub da_bnlh: Tensor<4>,
    pub rotation: RotationSeq,
    pub initial_state_bhpr: Tensor<4>,
    pub init_state_hpr: Option<Tensor<3>>,
}

Inputs of the physical-frame state-moments computation — the combined (γ + β) injections of one chunked Mamba-3 forward, plus the per-token cumulative rotation. Built at either pathway’s pre-SSD seam; C/D are never read. See the module header for the math.

Fields

xhat_bnlMhp: Tensor<6>

Combined value injections, already γ/β-scaled — M = 2·mimo_rank channels (γ channels first, then the shifted β channels). Zero at pads.

Shape

• [batch, nchunks, chunk_len, M, nheads, per_head_dim]

bhat_bnlMhr: Tensor<6>

Combined key injections: the rotated B̃ per γ channel and the shifted B̃ₜ₋₁ per β channel. Zero at pads.

Shape

• [batch, nchunks, chunk_len, M, nheads, state_rank]

da_bnlh: Tensor<4>

Log-decay Δₜ·Aₜ (negative; zero at pads ⇒ identity decay).

Shape

• [batch, nchunks, chunk_len, nheads]

rotation: RotationSeq

Per-token cumulative rotation over the padded sequence (sequence = nchunks·chunk_len; pad values are arbitrary — masked).

initial_state_bhpr: Tensor<4>

The SSM state entering this call (the cache’s ssm_bhpr — not the single-ssd kernel’s seed-augmented initial state; see module header).

Shape

• [batch, nheads, per_head_dim, state_rank]

init_state_hpr: Option<Tensor<3>>

Optional learnable initial state h₀, added to initial_state_bhpr exactly once.

Shape

• [nheads, per_head_dim, state_rank]

Implementations

impl Mamba3MomentsInput

pub fn detached(&self) -> Self

A value-identical copy detached from any autodiff graph (used by the diagnostic wrapper so the moments branch records no backward nodes).

pub fn state_moments_phys(&self, valid_len: usize) -> StateMoments

Exact pooled moments of every per-token physical-frame SSM state (see the module header). valid_len is the unpadded token count — states at zero-pad positions are excluded.

Returns raw sums; count = valid_len · per_head_dim samples per (batch, head) slice (MIMO ranks share the state, so the count is unchanged from Mamba-2).

pub fn state_moments_phys_recalculated(&self, valid_len: usize) -> StateMoments

Self::state_moments_phys with a custom recompute backward — the at-scale execution model. Mathematically identical (values and gradients, asserted by the tests); only the backward’s memory profile differs: plain autodiff retains every chunk’s states (the full per-token trajectory), while this node saves only the leaf inputs and re-materialises one chunk at a time during backprop (see moments/backward.rs).

The optional learnable init_state_hpr is folded into the initial state outside the node (a plain autodiff add), so its gradient flows through d_initial.