Module diag 

The same-step γ-correction on primitives — forward and analytic backward.

Same-step γ-correction on primitives — forward and analytic backward

Primitive (F) port of super::super::diag, plus the analytic backward the recompute node needs. The forward is used by super::serial_recalculated’s K5; the backward by super::combined_backward. Both carry the same SISO fast path: at mimo_rank == 1 the m × m Gram matrix is a scalar, so every matmul over the mimo_rank axis degenerates into a 1×K×1 / 1×1×N GEMM and is replaced by a reduction plus broadcast multiplies.

Forward (per (b, n, l, h)):

  qk_dot[m_out, m_in] = Σ_r C[m_out, r] · B[m_in, r]
  y_diag[m_out, p]    = γ · Σ_{m_in} qk_dot[m_out, m_in] · V[m_in, p]

Structs

DiagGrads

Gradients produced by y_diag_correction_backward — the y_diag term’s contribution to d_v, d_c, d_b and the whole of d_gamma (no other term consumes γ).

Functions

y_diag_correction

The γ-weighted same-step correction y_diag, on primitives.

y_diag_correction_backward

Analytic backward of y_diag_correction.

y_diag_correction_backward_mimo 🔒

General MIMO backward.

y_diag_correction_backward_siso 🔒

SISO (mimo_rank == 1) backward.

y_diag_correction_mimo 🔒

General MIMO y_diag on primitives.

y_diag_correction_siso 🔒

SISO (mimo_rank == 1) y_diag on primitives: a state_rank reduction plus a per-(b, n, l, h) scalar broadcast.