Module moments 

Closed-form per-token state moments (for state-PR diagnostics/penalties).

Exact per-token state moments from the chunkwise SSD (no state materialisation)

A token-by-token step loop exposes every SSM state hₜ ∈ ℝ^{per_head_dim × state_rank}; the chunkwise forward only materialises states at chunk boundaries. This module computes the pooled first/second moments of all per-token states — everything a participation-ratio diagnostic (or penalty) needs — in closed form, from tensors the SSD decomposition already builds, without ever materialising a [batch, sequence, nheads, per_head_dim, state_rank] state tensor.

Derivation

Within chunk n (per (batch, head), positions t ∈ [0, chunk_len)), unrolling the recurrence hₜ = Āₜhₜ₋₁ + xₜ ⊗ B̄ₜ from the state h₋ entering the chunk gives

  hₜ = dₜ·h₋ + sₜ ,   dₜ = exp(a_cum[t]) ,   sₜ = Σ_{j≤t} L[t,j] · xⱼ ⊗ B̄ⱼ

with L[t,j] = exp(a_cum[t] − a_cum[j]) — exactly the 1-semiseparable mask of SSD Step 1. Summing hₜᵀhₜ over t (the ᵀ contraction pools per_head_dim) splits into three chunk-level contractions:

  Σₜ hₜᵀhₜ = (Σₜ dₜ²)·h₋ᵀh₋                      (carried-state term)
           + h₋ᵀP + Pᵀh₋ ,  P = Xᵀ·diag(w)·B̄ ,  wⱼ = Σₜ dₜ·L[t,j]
           + B̄ᵀ·(K ∘ XXᵀ)·B̄ ,                    K = LᵀL

and the first moment likewise: Σₜ hₜ = (Σₜ dₜ)·h₋ + Xᵀ·diag(u)·B̄, uⱼ = Σₜ L[t,j]. Everything is a batched GEMM in the same size class as SSD Step 1 ([chunk_len, chunk_len] / [state_rank, state_rank] intermediates). The boundary states h₋ are recomputed with SSD Steps 2–3, so this function is pathway-agnostic (independent of which SSD variant produced y) and fully differentiable through plain autodiff.

Zero-pad positions (Δ=0 ⇒ Ā=1, B̄=0) carry the state unchanged, so counting them would replicate the final state into the moments; a validity mask excludes them from every Σₜ.