Struct StateMoments 

pub struct StateMoments {
    pub m2_bhrr: Tensor<4>,
    pub m1_bhr: Tensor<3>,
    pub count: usize,
}

Raw (un-normalised) first/second moments of the per-token SSM states of one block’s forward pass, pooled over tokens and per_head_dim rows.

Produced by forward_with_state_moments — closed-form for Mamba-2 (Mamba2SsdInput::state_moments, no state materialisation) and serial chunkwise for Mamba-3 (Mamba3MomentsInput::state_moments_phys, the per-token physical-frame states of the complex SSM).

Fields

m2_bhrr: Tensor<4>

Second-moment (Gram) sum Σₜ hₜᵀ hₜ — the ᵀ contraction pools the per_head_dim rows, the Σₜ the (unpadded) tokens.

Shape

• [batch, nheads, state_rank, state_rank]

m1_bhr: Tensor<3>

First-moment sum Σₜ Σₚ hₜ[p, :].

Shape

• [batch, nheads, state_rank]

count: usize

Samples pooled into each (batch, head) slice: valid_tokens · per_head_dim (grows additively under Self::merge).

Implementations

impl StateMoments

pub fn merge(self, other: Self) -> Self

Pool two moment sets (e.g. consecutive streamed forward calls, or separate eval batches). PR of the merged moments is the exact PR of the union of samples.

pub fn pool_batch(self) -> Self

Fold the batch dimension into the samples (batch-pooled moments with batch = 1), matching diagnostics that treat every (token, batch, per_head_dim) triple as one sample.

pub fn pr(&self, center: bool) -> Tensor<2>

Participation ratio (tr Σ)² / tr(Σ²) of the sample covariance, per (batch, head) slice; center subtracts the sample mean (Σ becomes the centered covariance instead of the raw second moment).

Differentiable (two traces, no eigendecomposition).

Shape

• output: [batch, nheads]

pub fn pr_complex(&self, pairing: &StatePairing, center: bool) -> Tensor<2>

Participation ratio of the Hermitian sample covariance, treating the state_rank axis as realified complex (or quaternionic) coordinates per pairing — the Mamba-3 counterpart of Self::pr.

With the pairing’s complex view c = x + iy, the Hermitian moment is M = A + iS with A = Σ(xxᵀ + yyᵀ) and S = Σ(xyᵀ − yxᵀ) — both linear recombinations of m2_bhrr sub-blocks, so centering Σ centers M identically. PR_ℂ = (tr M)² / tr(M²); the trace is real and equals the full real trace (frame-invariant), while tr(M²) = Σ|M_ab|² (= ‖A‖²_F + ‖S‖²_F for the fully-rotated complex case). One complex (or quaternionic) direction counts as one — the ×2 (×4) realified count is a representation artifact — so a rank-1 rotating conveyor reads PR_ℂ ≡ 1 where Self::pr reads up to the block size.

Un-rotated coordinates (partial rope_fraction) stay a real block U of the mixed Hermitian [[M, X], [Xᴴ, U]]; its trace and Σ|·|² join the sums (tr(·²) gains 2‖X‖²_F + ‖U‖²_F). Quaternionic pairing uses the same formulas with M_jk = Σ q̄ⱼqₖ (diagonal real, 4-component norms).

Differentiable; numerics mirror Self::pr (detached trace normalisation, min_positive / div_eps floors — see the comments there for why).

Shape

• output: [batch, nheads]

pub fn trace(&self) -> Tensor<2>

Raw uncentered state magnitude tr Σ = trace(m2)/count per (batch, head) — the mean squared state magnitude ⟨‖h‖²⟩, which is Self::pr’s numerator scale. Reported alongside PR to tell a genuine rank-1 state (PR → 1, magnitude healthy) apart from a state collapsing toward zero (where pr’s 1e-12 denominator clamp drags the ratio below its true floor of 1).

Shape

• output: [batch, nheads]

Trait Implementations

impl Clone for StateMoments

fn clone(&self) -> StateMoments

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for StateMoments

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.