burn_mamba/mamba3/rotation/

mod.rs

//! # Quaternion (k=4) rotational state — the non-abelian generalisation of RoPE
//!
//! Mamba-3's data-dependent RoPE realises a **complex-valued** SSM: the state
//! transition factors as a per-head scalar decay times a block-diagonal of
//! `2×2` rotations (paper Prop. *Complex-to-Real SSM Equivalence*), and because
//! `SO(2) ≅ U(1)` is **abelian** the cumulative rotation collapses to a
//! `cumsum` of angles and is absorbed into `B`/`C` (the "RoPE trick", Prop.
//! *Complex SSM, Data-Dependent RoPE Equivalence*).  See
//! [`crate::mamba3::rotation::rope::apply_rope`].
//!
//! This module implements the next rung of the ladder: a **quaternion**
//! (`k = 4`) rotational state, i.e. the transition's rotation lives in the
//! left-isoclinic subgroup `SU(2) ⊂ SO(4)` instead of `SO(2)`.  Unit
//! quaternions under multiplication are `SU(2)`, which is **non-abelian** and
//! contains non-solvable finite subgroups (the binary icosahedral group
//! `2I = SL(2,5)`, a double cover of `A₅`).  By Barrington's theorem this lifts
//! the layer's reachable state-tracking from the solvable/`TC⁰` regime (parity,
//! mod-k) toward `NC¹`, which abelian rotations provably cannot reach.
//!
//! ## What survives, what changes
//!
//! The key fact (derivable purely from telescoping + orthogonality, **without**
//! commutativity — see the crate discussion) is that the RoPE *factoring*
//! survives intact: with the **ordered** cumulative rotation
//! `Pₜ = Rₜ Rₜ₋₁ ⋯ R₁`,
//!
//! ```text
//!   Cₜᵀ (Rₜ⋯Rᵢ₊₁) Bᵢ  =  (Pₜᵀ Cₜ)ᵀ (Pᵢᵀ Bᵢ)  =  C̄ₜᵀ B̄ᵢ ,
//! ```
//!
//! so the scalar-decay SSD core (`L ⊙ C̄B̄ᵀ`) is **unchanged** — only the
//! projections `B̄ᵢ = Pᵢᵀ Bᵢ`, `C̄ₜ = Pₜᵀ Cₜ` are rotated.  What is lost is the
//! closed-form `cumsum`: the cumulative rotation must be built by an
//! **associative scan over the per-step quaternions** ([`quat_cumprod`]) rather
//! than a sum of angles.  Because a product of unit quaternions is again a unit
//! quaternion, the scan stays exactly orthogonal (no drift, no `wrap_angle`
//! needed), and the cross-chunk carry is a single quaternion per block/head —
//! the exact analogue of `cum_angle` in the existing caches.
//!
//! `SO(2)` (today's `apply_rope`) is the abelian collapse: restricting each
//! quaternion to a single fixed axis makes them commute and reduces
//! [`quat_cumprod`] to a `cumsum` of half-angles (asserted in the tests).
//!
//! ## Pipeline (the `k = 4` instantiation of the rotation block)
//!
//! ```text
//!   per-step unit quaternion qₜ      (materialise from the in-projection; caller)
//!        │  quat_cumprod (assoc. scan, + cross-chunk carry)
//!        ▼
//!   cumulative rotation Qₜ
//!        │  rotate_state_rank_blocks(B, conj(Qₜ)) , rotate_state_rank_blocks(C, conj(Qₜ))
//!        ▼
//!   B̄, C̄  ──►  standard scalar-decay SSD  (unchanged)
//! ```
//!
//! For [`RotationKind::Rotor4D`] the same pipeline runs with the two factors
//! stacked along the block axis, and the last step becomes the two-sided
//! `rotate_state_rank_blocks_two_sided(B, conj(Qₜ), Tₜ)`.
//!
//! ## `SO(4)`: the whole rotation group of a block ([`RotationKind::Rotor4D`])
//!
//! `SU(2)` is only half of what a 4-block can turn by. The general element of
//! `SO(4) ≅ (SU(2)×SU(2))/±1` is the **two-sided** product
//!
//! ```text
//!   Rₜ(v) = qₜ ⊗ v ⊗ p̄ₜ            (a rotor; left factor q, right factor p)
//! ```
//!
//! and everything above survives it, because the factoring never used more
//! than "the per-step maps compose, and each is orthogonal". Composing,
//!
//! ```text
//!   Pₜ(v) = Qₜ ⊗ v ⊗ T̄ₜ ,   Qₜ = qₜ⊗⋯⊗q₁ ,  Tₜ = pₜ⊗⋯⊗p₁
//!   Pₜ⁻¹(v) = Qₜ* ⊗ v ⊗ Tₜ    ⇒   B̄ᵢ = Qᵢ* ⊗ Bᵢ ⊗ Tᵢ ,  C̄ₜ likewise
//! ```
//!
//! — note the conjugation reverses the right-hand order **twice**, so `T`
//! accumulates by the *same* left fold as `Q`: one [`quat_cumprod`] over a
//! doubled block axis, not a second, reversed scan. The cost over
//! [`Quaternion4D`](RotationKind::Quaternion4D) is twice the generator
//! channels, twice the scan's block axis, and one extra [`quat_mul`] per
//! `B`/`C` application; the SSD core is still untouched.
//!
//! Why bother, when left and right factors *commute* with each other and so add
//! no "more non-abelianness": left multiplication is **isoclinic** — `L_q` turns
//! both invariant planes of the block by the same angle — so `SU(2)` cannot
//! produce two independent plane angles, and in particular does not contain the
//! abelian `SO(2)²` rotation it was introduced to generalise. The right factor
//! is exactly what opens the maximal torus (plane angles `a−b`, `a+b` for the
//! half-angles of `q`, `p`), making the ladder
//! `Complex2D ⊂ Rotor4D ⊃ Quaternion4D` a real one. It also contains the
//! adjoint action `v ↦ q ⊗ v ⊗ q̄`, i.e. a faithful `SO(3)` on the block's
//! imaginary part — so a group like `A₅` can be tracked as itself rather than
//! through its double cover `2I`, where `±g` denote one element but two
//! different states. And as a representation of `SU(2)×SU(2)`, `ℍ` is the
//! irreducible tensor product `(½,½)`, not `(½,0) ⊕ (0,½)`: no arrangement of
//! left-only blocks reproduces it.
//!
//! `SO(4)` is the ceiling for `k = 4` (the largest norm-preserving transition
//! group of a block), and `k = 4` is the last rung with a cheap closed form —
//! at `k = 8` the octonions are non-associative and the scan itself breaks.
//!
//! ## The bottom rung: [`RotationKind::Real1D`]
//!
//! `k = 1`, the trivial group: no rotation, a purely real transition. It is the
//! ablation the ladder is measured against, and it is *structural* rather than a
//! zeroed knob — the in-projection spends no channels on rotation, the cache
//! carries no accumulator ([`RotationState::Real`]), and `B`/`C` reach the SSD
//! core untouched. Switching the rotation off is therefore a choice of *kind*;
//! `rope_fraction` only ever narrows a rotation that exists.
//!
//! Quaternion layout: the last axis has size 4 and holds `(w, x, y, z)` with
//! `w` the real part.  A `state_rank` of `r = 4·J` is treated as `J` independent
//! quaternion blocks; the rotation acts within each block, exactly as RoPE acts
//! within each `2`-pair.  [`Mamba3`](crate::mamba3::mamba3::Mamba3) selects it
//! with [`RotationKind::Quaternion4D`] and drives it through one
//! [`RotationSpec`] (the SSD kernels themselves need no edits).

/// Rotary (RoPE) application: the mechanical pairwise rotation the abelian
/// pathway factors into B/C.
pub mod rope;

use crate::mamba3::rotation::rope::{apply_rope_partial, wrap_angle};
use burn::module::Module;
use burn::prelude::*;

// ---------------------------------------------------------------------------
// Rotation kind (config switch) and cache accumulator variant
// ---------------------------------------------------------------------------

/// Which rotational-state algebra the block uses for the data-dependent
/// transition rotation absorbed into `B`/`C`.
///
/// - [`Real1D`](RotationKind::Real1D) — the trivial group: no rotation at all.
///   The transition is the plain scalar decay, i.e. a **real** SSM (Mamba-2's
///   transition under Mamba-3's trapezoid). The in-projection spends *no*
///   channels on rotation and the cache carries no accumulator
///   ([`RotationState::Real`]) — this is the rotation ablation, which is why
///   `rope_fraction` has no `0.0` setting.
/// - [`Complex2D`](RotationKind::Complex2D) — the abelian `SO(2)`/complex RoPE
///   that Mamba-3 ships: cumulative *angles* via `cumsum`, applied by
///   [`apply_rope`](crate::mamba3::rotation::rope::apply_rope). The default; behaviourally unchanged.
/// - [`Quaternion4D`](RotationKind::Quaternion4D) — the non-abelian
///   `SU(2) ⊂ SO(4)` quaternion rotation of this module: cumulative *product*
///   via [`quat_cumprod`], applied by [`rotate_state_rank_blocks`]. Richer
///   state-tracking; selects the [`RotationState::Quaternion`] cache accumulator.
/// - [`Rotor4D`](RotationKind::Rotor4D) — the **whole** rotation group of a
///   4-block, `SO(4) ≅ (SU(2)×SU(2))/±1`: a *two-sided* quaternion product
///   `v ↦ q ⊗ v ⊗ p̄`. Two per-step quaternions instead of one (twice the
///   generator channels, one scan over a doubled block axis); selects the
///   [`RotationState::Rotor`] accumulator.
///
/// The kinds are a ladder — `Real1D ⊂ Complex2D ⊂ Rotor4D` and
/// `Real1D ⊂ Quaternion4D ⊂ Rotor4D` — but the middle two are
/// **incomparable**, which is why `Rotor4D` exists.
/// Left multiplication is *isoclinic*: `L_q` turns both invariant planes of the
/// block by the same angle (`L_i` sends `1↦i` **and** `j↦k`), so
/// `Quaternion4D` cannot express two independent per-pair angles — it does not
/// contain the abelian rotation it generalises. The right factor opens the full
/// maximal torus (plane angles `a−b` and `a+b` for half-angles `a`, `b`), and
/// with it every element of `SO(4)`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
pub enum RotationKind {
    /// No rotation: a real transition (the trivial group). No rotation
    /// channels, no cumulative accumulator, `B`/`C` pass through untouched.
    Real1D,
    /// Abelian complex (`SO(2)`) RoPE — the current default behaviour.
    #[default]
    Complex2D,
    /// Non-abelian quaternion (`SU(2)`, left-isoclinic) rotation.
    Quaternion4D,
    /// The full `SO(4)` rotation of each 4-block: two-sided `v ↦ q ⊗ v ⊗ p̄`.
    Rotor4D,
}

impl RotationKind {
    /// How many cumulative-rotation quaternions the accumulator carries per
    /// (head, 4-block): one for [`Quaternion4D`](RotationKind::Quaternion4D),
    /// **two** for [`Rotor4D`](RotationKind::Rotor4D) (the left and right
    /// factors). Meaningless for the non-quaternion kinds
    /// ([`Complex2D`](RotationKind::Complex2D) accumulates angles,
    /// [`Real1D`](RotationKind::Real1D) accumulates nothing) — reported as 1.
    ///
    /// The factors are stacked along one block axis, so every quaternion
    /// primitive — the generator split, [`quat_from_scaled_axis`], the
    /// [`quat_cumprod`] scan, [`quat_normalize`] — runs **once** over
    /// `quat_factors · blocks` blocks and needs no `Rotor4D` branch of its own.
    /// Only the *application* to `B`/`C` differs.
    pub fn quat_factors(self) -> usize {
        match self {
            RotationKind::Rotor4D => 2,
            RotationKind::Real1D | RotationKind::Complex2D | RotationKind::Quaternion4D => 1,
        }
    }
}

/// Everything the rotation needs from the block: which algebra, how much of
/// `state_rank` it turns, and how far one step may turn it.
///
/// Carried by [`Mamba3`](crate::mamba3::mamba3::Mamba3) and handed to
/// [`rotate_bc_forward`] / [`rotate_bc_step`] (and to the constant-input
/// shortcut) so that every site derives the per-step rotation from **one**
/// definition — the three used to spell the formula out separately, which is
/// exactly the kind of duplication that lets `forward`, `step` and
/// `step_infinite` drift apart.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct RotationSpec {
    /// Which rotational-state algebra ([`RotationKind`]).
    pub kind: RotationKind,
    /// How many leading `state_rank` entries are rotated (see
    /// [`Mamba3Config::rope_fraction`](crate::mamba3::mamba3::Mamba3Config::rope_fraction)).
    /// For [`RotationKind::Quaternion4D`] this is a multiple of 4; for
    /// [`RotationKind::Real1D`] it is `0` and nothing reads it.
    pub rope_dim: usize,
    /// The per-step rotation bound in half-turns per unit `Δ`
    /// ([`Mamba3Config::rotation_range`](crate::mamba3::mamba3::Mamba3Config::rotation_range)):
    /// one step turns by at most `range · π · Δ`.
    pub range: f64,
}

/// The stateless payload of [`RotationState::Real`] — a [`Module`] holding no
/// tensors, so a real transition's cache slot allocates nothing and converts
/// between backends by `Clone`.
#[derive(Module, Debug, Default)]
pub struct NoRotation;

/// The cumulative-rotation accumulator carried between calls in a Mamba-3 cache
/// — the variant matching the block's [`RotationKind`].
///
/// - [`Real`](RotationState::Real) — nothing at all: a real transition has no
///   cumulative rotation to carry.
/// - [`Angle`](RotationState::Angle) — abelian per-pair cumulative RoPE angle,
///   shape `[batch, nheads, num_rope_angles]` (today's `cum_angle`).
/// - [`Quaternion`](RotationState::Quaternion) — per-block cumulative unit
///   quaternion, shape `[batch, nheads, blocks, 4]`, produced by
///   [`quat_cumprod`].
///
/// This is the cache-level counterpart of [`RotationKind`]. It is defined here
/// (the rotation module owns the accumulator type); substituting it for the
/// pathway caches' `cum_angle_bha` field happens together with the forward/step
/// wiring that consumes it.
#[derive(Module, Debug)]
pub enum RotationState {
    /// [`RotationKind::Real1D`]'s empty accumulator: a real transition composes
    /// nothing between calls. It carries a [`NoRotation`] only because Burn's
    /// `Module` derive takes exactly one field per enum variant.
    Real(NoRotation),
    /// Abelian RoPE cumulative angle, shape `[batch, nheads, num_rope_angles]`.
    Angle(Tensor<3>),
    /// Quaternion cumulative rotation, shape `[batch, nheads, blocks, 4]`.
    Quaternion(Tensor<4>),
    /// `SO(4)` cumulative rotation, shape `[batch, nheads, 2·blocks, 4]`: the
    /// left factors `Qₜ = qₜ⊗⋯⊗q₁` in the first `blocks` entries of the block
    /// axis, the right factors `Tₜ = pₜ⊗⋯⊗p₁` in the second (see
    /// [`split_rotor`]).
    ///
    /// One tensor rather than two so the scan, the normalisation and the cache
    /// plumbing stay single-call; the conjugation in `v ↦ q v p̄` reverses the
    /// right-hand order **twice**, so `T` accumulates with the very same
    /// left-fold as `Q` and no reversed scan is needed.
    Rotor(Tensor<4>),
}

impl RotationState {
    /// The empty accumulator of [`RotationKind::Real1D`].
    pub fn real() -> Self {
        RotationState::Real(NoRotation)
    }

    /// Zero-initialised abelian angle accumulator `[batch, nheads, num_rope_angles]`.
    pub fn zeros_angle(
        batch: usize,
        nheads: usize,
        num_rope_angles: usize,
        device: &Device,
    ) -> Self {
        RotationState::Angle(Tensor::zeros([batch, nheads, num_rope_angles], device))
    }

    /// Identity-initialised quaternion accumulator `[batch, nheads, blocks, 4]`
    /// (every block is the identity quaternion `(1, 0, 0, 0)`).
    pub fn identity_quaternion(
        batch: usize,
        nheads: usize,
        blocks: usize,
        device: &Device,
    ) -> Self {
        let w = Tensor::ones([batch, nheads, blocks, 1], device);
        let xyz = Tensor::zeros([batch, nheads, blocks, 3], device);
        RotationState::Quaternion(Tensor::cat(vec![w, xyz], 3))
    }

    /// The variant's name, for assertion messages.
    fn variant(&self) -> &'static str {
        match self {
            RotationState::Real(_) => "Real",
            RotationState::Angle(_) => "Angle",
            RotationState::Quaternion(_) => "Quaternion",
            RotationState::Rotor(_) => "Rotor",
        }
    }

    /// Check this is the empty [`Real`](RotationState::Real) accumulator and
    /// hand it back; panics on any other variant, so a cache built for a
    /// rotating kind cannot be fed to a [`Real1D`](RotationKind::Real1D) block.
    pub fn expect_real(self) -> Self {
        match self {
            RotationState::Real(_) => self,
            other => panic!("RotationState is {}, expected Real", other.variant()),
        }
    }

    /// Unwrap the abelian angle accumulator; panics if this is a quaternion.
    pub fn angle(self) -> Tensor<3> {
        match self {
            RotationState::Angle(a) => a,
            other => panic!("RotationState is {}, expected Angle", other.variant()),
        }
    }

    /// Identity-initialised `SO(4)` accumulator `[batch, nheads, 2·blocks, 4]`
    /// (both factors of every block are the identity quaternion `(1, 0, 0, 0)`).
    pub fn identity_rotor(batch: usize, nheads: usize, blocks: usize, device: &Device) -> Self {
        let w = Tensor::ones([batch, nheads, 2 * blocks, 1], device);
        let xyz = Tensor::zeros([batch, nheads, 2 * blocks, 3], device);
        RotationState::Rotor(Tensor::cat(vec![w, xyz], 3))
    }

    /// The identity accumulator for `kind` — the one place a fresh cache's
    /// rotation state is built, for every pathway and both cache types.
    pub fn identity(
        kind: RotationKind,
        batch: usize,
        nheads: usize,
        num_rope_angles: usize,
        num_quat_blocks: usize,
        device: &Device,
    ) -> Self {
        match kind {
            RotationKind::Real1D => RotationState::real(),
            RotationKind::Complex2D => {
                RotationState::zeros_angle(batch, nheads, num_rope_angles, device)
            }
            RotationKind::Quaternion4D => {
                RotationState::identity_quaternion(batch, nheads, num_quat_blocks, device)
            }
            RotationKind::Rotor4D => {
                RotationState::identity_rotor(batch, nheads, num_quat_blocks, device)
            }
        }
    }

    /// Unwrap the quaternion accumulator; panics if this is an angle.
    pub fn quaternion(self) -> Tensor<4> {
        match self {
            RotationState::Quaternion(q) => q,
            other => panic!("RotationState is {}, expected Quaternion", other.variant()),
        }
    }

    /// Unwrap the `SO(4)` accumulator (`[batch, nheads, 2·blocks, 4]`, both
    /// factors stacked); panics for any other variant.
    pub fn rotor(self) -> Tensor<4> {
        match self {
            RotationState::Rotor(q) => q,
            other => panic!("RotationState is {}, expected Rotor", other.variant()),
        }
    }

    /// The stacked quaternion accumulator for `kind`, together with the number
    /// of **state** 4-blocks it covers — half its block axis for
    /// [`RotationKind::Rotor4D`], which stacks two factors there.
    ///
    /// Panics on a variant the kind does not use: the two are the same rank and
    /// differ only in the length of one axis, so a mismatched cache would
    /// otherwise be reinterpreted rather than rejected.
    fn quat_stack(self, kind: RotationKind) -> (Tensor<4>, usize) {
        match (kind, self) {
            (RotationKind::Quaternion4D, RotationState::Quaternion(q)) => {
                let blocks = q.dims()[2];
                (q, blocks)
            }
            (RotationKind::Rotor4D, RotationState::Rotor(q)) => {
                let stacked = q.dims()[2];
                assert_eq!(stacked % 2, 0, "a Rotor accumulator stacks two factors");
                (q, stacked / 2)
            }
            (kind, other) => panic!(
                "RotationState is {}, which is not {kind:?}'s accumulator",
                other.variant()
            ),
        }
    }

    /// Run the [`NaN`/`Inf` guards](burn_stack::modules::misc::sanity) on the held tensor.
    pub fn sanity(&self) {
        match self {
            RotationState::Real(_) => {}
            RotationState::Angle(a) => burn_stack::modules::sanity(a),
            RotationState::Quaternion(q) | RotationState::Rotor(q) => burn_stack::modules::sanity(q),
        }
    }
}

// ---------------------------------------------------------------------------
// Quaternion algebra on the trailing `(w, x, y, z)` axis
// ---------------------------------------------------------------------------

/// Hamilton product `a ⊗ b` of two quaternion tensors.
///
/// Both inputs have shape `[..., 4]` with the last axis ordered `(w, x, y, z)`;
/// the product is computed component-wise and broadcasts over the leading dims.
/// Quaternion multiplication is **non-commutative** (`a ⊗ b ≠ b ⊗ a` in
/// general) but associative.
///
/// Identifying `ℝ⁴` with the quaternions, left-multiplication `v ↦ a ⊗ v` is
/// exactly the action of the `4×4` rotation matrix [`quat_to_rot4`]`(a)`, so
/// this is also how a rotation is *applied* to a state/`B`/`C` block (see
/// [`rotate_state_rank_blocks`]).
pub fn quat_mul<const D: usize>(a: Tensor<D>, b: Tensor<D>) -> Tensor<D> {
    let n = D - 1;
    let aw = a.clone().narrow(n, 0, 1);
    let ax = a.clone().narrow(n, 1, 1);
    let ay = a.clone().narrow(n, 2, 1);
    let az = a.narrow(n, 3, 1);
    let bw = b.clone().narrow(n, 0, 1);
    let bx = b.clone().narrow(n, 1, 1);
    let by = b.clone().narrow(n, 2, 1);
    let bz = b.narrow(n, 3, 1);

    // Hamilton product (each term is shape [..., 1]).
    let w = aw.clone() * bw.clone()
        - ax.clone() * bx.clone()
        - ay.clone() * by.clone()
        - az.clone() * bz.clone();
    let x = aw.clone() * bx.clone() + ax.clone() * bw.clone() + ay.clone() * bz.clone()
        - az.clone() * by.clone();
    let y = aw.clone() * by.clone() - ax.clone() * bz.clone()
        + ay.clone() * bw.clone()
        + az.clone() * bx.clone();
    let z = aw * bz + ax * by - ay * bx + az * bw;

    Tensor::cat(vec![w, x, y, z], n)
}

/// Quaternion conjugate `q* = (w, −x, −y, −z)` (shape `[..., 4]`).
///
/// For a **unit** quaternion `q* = q⁻¹`, and the corresponding rotation matrix
/// satisfies `Lₚ⋆ = Lₚᵀ = Lₚ⁻¹`.  Hence rotating by the *inverse* cumulative
/// rotation (`B̄ = Pᵀ B`) is `rotate_state_rank_blocks(B, conj(Q))`.
pub fn quat_conj<const D: usize>(q: Tensor<D>) -> Tensor<D> {
    let n = D - 1;
    let w = q.clone().narrow(n, 0, 1);
    let xyz = q.narrow(n, 1, 3);
    Tensor::cat(vec![w, -xyz], n)
}

/// Normalise quaternions to unit norm along the last axis (shape `[..., 4]`).
///
/// The per-step rotation is materialised from a raw, unconstrained projection
/// and normalised here so it is a genuine unit quaternion (an element of
/// `SU(2)`), the analogue of `tanh(θ)·π` bounding the RoPE angle.  A tiny floor
/// guards the zero-quaternion.
pub fn quat_normalize<const D: usize>(q: Tensor<D>) -> Tensor<D> {
    let n = D - 1;
    // Clamp the sum-of-squares *before* `sqrt`: at a zero quaternion the forward
    // `sqrt(0)=0` is fine, but `sqrt`'s backward is `1/(2·0)=∞`, and `∞·(2·0)=NaN`.
    // Clamping pre-`sqrt` puts the degenerate point in `clamp_min`'s flat region,
    // so its gradient is a finite 0 (and a genuine unit quaternion, sumsq=1, is
    // untouched). The floor also keeps `norm` away from 0 for the division.
    //
    // The floor is the dtype-aware `div_eps` applied to the *sum-of-squares*
    // (giving a norm floor of `√div_eps`). It must engage as a representable
    // normal in the working dtype: in f16 a `div_eps²`-sized floor (~5e-7) would
    // underflow below the min-normal (~6.1e-5) and silently no-op, so we floor
    // the squared quantity at `div_eps` itself, which sits above each format's
    // denormal floor by construction.
    let eps = burn_stack::utils::div_eps(q.dtype());
    let norm = (q.clone() * q.clone()).sum_dim(n).clamp_min(eps).sqrt();
    q / norm
}

/// Euclidean norm over the last axis, formed **scale-free** so it cannot
/// overflow: the components are divided by their (detached) largest magnitude
/// before squaring, and the result is scaled back.
///
/// Squaring the raw components is the obvious way and the wrong one here,
/// because the inputs are raw in-projection channels: `‖r‖²` overflows f32 at
/// `|r| ≈ 2e19` and **f16 at `|r| ≈ 250`**, which is an ordinary activation.
/// The overflow does not announce itself — `∞` divides back to `0`, so a very
/// large generator would silently produce *no* rotation, the exact opposite of
/// the intended "turn as far as the bound allows". Same trick, same reason, as
/// [`RmsNorm`](burn_stack::modules::RmsNorm)'s fp16 path.
///
/// The sum of squares is floored by `div_eps` before the `sqrt`, so a zero
/// vector lands in `clamp_min`'s flat region and backprops to a finite `0`
/// instead of `sqrt`'s singular `1/(2·0)`.
///
/// # Shapes
/// - `t`  : `[..., n]`
/// - out  : `[..., 1]`
pub(crate) fn safe_norm<const D: usize>(t: Tensor<D>) -> Tensor<D> {
    let n = D - 1;
    let eps = burn_stack::utils::div_eps(t.dtype());
    // Detached: the rescaling is a numerical device, not part of the function
    // being differentiated (`d‖t‖/dt = t̂` either way).
    let scale = t.clone().abs().max_dim(n).detach() + eps; // [..., 1]
    let unit = t / scale.clone(); // components ≤ 1
    (unit.clone() * unit).sum_dim(n).clamp_min(eps).sqrt() * scale
}

/// Materialise a unit quaternion from a **scaled rotation vector** `g ∈ ℝ³`
/// (axis · angle) via the exponential map — the data-dependent "materialise
/// `Rₜ`" step, analogous to RoPE's `Δₜ · π · tanh(θₜ)` angle.
///
/// With `‖g‖ = angle` and `ĝ = g / angle` the axis, returns the unit quaternion
/// `q = (cos(angle/2), sin(angle/2)·ĝ)`.  A vanishing `g` maps to the identity
/// `(1, 0, 0, 0)`, so scaling `g` by a small `Δₜ` (the discretisation step)
/// yields a near-identity rotation — exactly the regime where a small step
/// barely rotates the state.  The `sin(angle/2)/angle` factor is the numerically
/// stable form of the (otherwise `0/0`) per-component scale near `g = 0`.
///
/// # Shapes
/// - `g` : `[..., 3]`
/// - out : `[..., 4]` (ordered `(w, x, y, z)`), unit norm.
pub fn quat_from_scaled_axis<const D: usize>(g: Tensor<D>) -> Tensor<D> {
    let n = D - 1;
    // `safe_norm` both guards the origin (a zero generator would otherwise hit
    // `sqrt`'s singular backward and yield NaN — the FiLM-triggered decoder
    // NaN) and keeps the sum of squares from overflowing at large `g`.
    let angle = safe_norm(g.clone()); // [..., 1]
    let half = angle.clone() * 0.5;
    let w = half.clone().cos(); // [..., 1]
    // sin(angle/2) / angle  → 1/2 as angle → 0 (no rotation); `angle ≥ √div_eps`
    // after the pre-`sqrt` clamp above, so the division is already guarded.
    let scale = half.sin() / angle; // [..., 1]
    let v = g * scale; // [..., 3]
    quat_normalize(Tensor::cat(vec![w, v], n))
}

/// Bound a rotation vector's **magnitude**, leaving its direction alone:
/// returns `max_angle · tanh(‖r‖) · r̂`.
///
/// This is the quaternion counterpart of the abelian path's `π·tanh(ϑ)`, and
/// the difference is deliberate. Squashing each of the three raw channels
/// *separately* would bound the rotation vector to a **cube**: the reachable
/// angle would depend on the axis (`max_angle` about a coordinate axis but
/// `√3·max_angle` about the diagonal), and — worse — `tanh` applied per
/// component moves the *direction* too, so the axis a given projection selects
/// would depend on how large the projection is. Bounding the norm keeps the
/// axis exactly `r̂` and the angle a function of `‖r‖` alone, so axis and angle
/// are independent knobs.
///
/// Near `r = 0` the map is `≈ max_angle · r` (the `tanh(n)/n → 1` limit); the
/// `sum-of-squares` floor is the same pre-`sqrt` clamp
/// [`quat_from_scaled_axis`] uses, and puts the degenerate point in the flat
/// region of `clamp_min` with the correct finite gradient.
///
/// # Shapes
/// - `r`  : `[..., 3]`
/// - out  : `[..., 3]`, with `‖out‖ ≤ max_angle` (attained once `tanh` saturates).
pub fn bound_rotation_vector<const D: usize>(r: Tensor<D>, max_angle: f64) -> Tensor<D> {
    let norm = safe_norm(r.clone()); // [..., 1]
    let scale = norm.clone().tanh() * max_angle / norm;
    r * scale
}

/// Materialise the `4×4` orthogonal matrix of left-multiplication by `q`.
///
/// Maps `q` of shape `[..., 4]` to `[..., 4, 4]` such that, for `v` of shape
/// `[..., 4]`, `Lq · v == quat_mul(q, v)`.  Concretely (rows = output coords,
/// cols = input coords, all in `(w, x, y, z)` order):
///
/// ```text
///   ⎡ w  -x  -y  -z ⎤
///   ⎢ x   w  -z   y ⎥
///   ⎢ y   z   w  -x ⎥
///   ⎣ z  -y   x   w ⎦
/// ```
///
/// For a unit `q` this is orthogonal with `det = 1` (a left-isoclinic rotation).
/// Provided mainly for the generic / verification path; the cheap way to apply a
/// rotation is [`rotate_state_rank_blocks`] (a quaternion product, no `4×4`
/// materialisation).  `DR` must equal `D + 1`.
pub fn quat_to_rot4<const D: usize, const DR: usize>(q: Tensor<D>) -> Tensor<DR> {
    assert_eq!(D + 1, DR, "quat_to_rot4 maps rank D to rank D+1");
    let n = D - 1;
    let w = q.clone().narrow(n, 0, 1);
    let x = q.clone().narrow(n, 1, 1);
    let y = q.clone().narrow(n, 2, 1);
    let z = q.narrow(n, 3, 1);

    // Each row is a [..., 4] tensor (the four column entries).
    let row0 = Tensor::cat(vec![w.clone(), -x.clone(), -y.clone(), -z.clone()], n);
    let row1 = Tensor::cat(vec![x.clone(), w.clone(), -z.clone(), y.clone()], n);
    let row2 = Tensor::cat(vec![y.clone(), z.clone(), w.clone(), -x.clone()], n);
    let row3 = Tensor::cat(vec![z, -y, x, w], n);

    // Stack the rows along a freshly inserted row axis → [..., 4, 4].
    Tensor::cat(
        vec![
            row0.unsqueeze_dim::<DR>(n),
            row1.unsqueeze_dim::<DR>(n),
            row2.unsqueeze_dim::<DR>(n),
            row3.unsqueeze_dim::<DR>(n),
        ],
        n,
    )
}

// ---------------------------------------------------------------------------
// Rotation application on the state_rank axis
// ---------------------------------------------------------------------------

/// Apply a per-block quaternion rotation to the `state_rank` axis of `v`.
///
/// `v` has shape `[..., state_rank]` with `state_rank = 4·J`, viewed as `J`
/// independent quaternion blocks; `q` has shape `[..., J, 4]` (one unit
/// quaternion per block, same leading dims as `v`).  Returns `q ⊗ v` per block,
/// i.e. the rotation `L_q` applied within each `4`-block, reshaped back to
/// `[..., state_rank]`.
///
/// This is the generalisation of RoPE's per-pair `2×2` rotation to per-block
/// `4×4`.  To rotate by the *inverse* cumulative rotation when absorbing into
/// `B`/`C` (`B̄ = Pᵀ B`), pass `q = conj(Qcum)`:
/// `rotate_state_rank_blocks(b, conj(qcum))`.
///
/// `DB` must equal `D + 1` (the block-split inserts the `J` axis).
pub fn rotate_state_rank_blocks<const D: usize, const DB: usize>(
    v: Tensor<D>,
    q: Tensor<DB>,
) -> Tensor<D> {
    assert_eq!(
        D + 1,
        DB,
        "rotate_state_rank_blocks splits one axis into (J, 4)"
    );
    let dims = v.dims();
    let state_rank = dims[D - 1];
    assert_eq!(
        state_rank % 4,
        0,
        "state_rank must be a multiple of 4 (quaternion blocks)"
    );
    let blocks = state_rank / 4;

    // Build the block-split shape [..., J, 4] (rank DB) and the flat shape
    // [..., state_rank] (rank D) for the round trip.
    let mut split_shape = [0usize; DB];
    split_shape[..D - 1].copy_from_slice(&dims[..D - 1]);
    split_shape[DB - 2] = blocks;
    split_shape[DB - 1] = 4;

    let v_blocks = v.reshape(split_shape); // [..., J, 4]
    let rotated = quat_mul(q, v_blocks); // L_q applied per block
    rotated.reshape(dims) // [..., state_rank]
}

/// Apply a per-block **two-sided** quaternion rotation to the `state_rank` axis
/// of `v`: `v ↦ ql ⊗ v ⊗ qr` per 4-block.
///
/// This is the general `SO(4)` element ([`RotationKind::Rotor4D`]); pass
/// `qr = (1,0,0,0)` to recover [`rotate_state_rank_blocks`]. Absorbing the
/// *inverse* cumulative rotation into `B`/`C` (`B̄ = P⁻¹B` with
/// `P(v) = Q v T̄`) is `ql = conj(Q)`, `qr = T` — note the conjugate is on the
/// **left** factor only.
///
/// Shapes as [`rotate_state_rank_blocks`]: `v` is `[..., state_rank]`, both
/// quaternions `[..., J, 4]`, and `DB = D + 1`.
pub fn rotate_state_rank_blocks_two_sided<const D: usize, const DB: usize>(
    v: Tensor<D>,
    ql: Tensor<DB>,
    qr: Tensor<DB>,
) -> Tensor<D> {
    assert_eq!(
        D + 1,
        DB,
        "rotate_state_rank_blocks_two_sided splits one axis into (J, 4)"
    );
    let dims = v.dims();
    let state_rank = dims[D - 1];
    assert_eq!(
        state_rank % 4,
        0,
        "state_rank must be a multiple of 4 (quaternion blocks)"
    );

    let mut split_shape = [0usize; DB];
    split_shape[..D - 1].copy_from_slice(&dims[..D - 1]);
    split_shape[DB - 2] = state_rank / 4;
    split_shape[DB - 1] = 4;

    let v_blocks = v.reshape(split_shape); // [..., J, 4]
    quat_mul(quat_mul(ql, v_blocks), qr).reshape(dims)
}

/// Split a stacked [`RotationState::Rotor`] accumulator `[..., 2·J, 4]` into
/// its `(left, right)` factors, each `[..., J, 4]`.
pub fn split_rotor<const D: usize>(q: Tensor<D>) -> (Tensor<D>, Tensor<D>) {
    let stacked = q.dims()[D - 2];
    assert_eq!(stacked % 2, 0, "a Rotor accumulator stacks two factors");
    let blocks = stacked / 2;
    (
        q.clone().narrow(D - 2, 0, blocks),
        q.narrow(D - 2, blocks, blocks),
    )
}

// ---------------------------------------------------------------------------
// Cumulative rotation scan (the associative, non-abelian replacement for cumsum)
// ---------------------------------------------------------------------------

/// Cumulative (ordered, left-accumulating) quaternion product along the
/// sequence axis, with a cross-chunk carry.
///
/// This is the non-abelian analogue of the cumulative *sum of angles* used by
/// RoPE: where complex rotations compose by adding angles (a `cumsum`),
/// quaternions compose by multiplication, which is order-dependent, so a real
/// scan is required.
///
/// # Shapes
/// - `q_bshj4` : `[batch, sequence, nheads, J, 4]` per-step **unit** quaternions
///   (block count `J = state_rank / 4`).
/// - `init`    : optional carry `[batch, nheads, J, 4]` — the cumulative
///   rotation at the end of the previous chunk (identity `(1,0,0,0)` for a fresh
///   start).
/// - returns `(cum, final_carry)` where `cum` is `[batch, sequence, nheads, J, 4]`
///   with `cum[:, t] = qₜ ⊗ qₜ₋₁ ⊗ ⋯ ⊗ q₀ ⊗ init` (newest on the left, matching
///   `Pₜ = Rₜ ⋯ R₁`), and `final_carry` `[batch, nheads, J, 4]` is `cum[:, −1]`
///   to thread into the next chunk.
///
/// Running this over a split sequence while threading `final_carry` is exactly
/// equal to running it over the whole sequence (asserted in the tests) — the
/// chunked-prefill / streaming guarantee, here for the rotation accumulator.
///
/// Implemented as a **Hillis–Steele** inclusive associative scan: the quaternion
/// product is associative (just not commutative), so a log-depth scan applies as
/// long as operand order is preserved (newest-on-left). Each doubling step is a
/// single full-tensor [`quat_mul`] plus a sequence shift, so the *sequential
/// dependency depth* is `O(log sequence)` rather than the `O(sequence)` of a
/// token-by-token loop — the same values, but a handful of large batched kernels
/// instead of thousands of serialized tiny ones (and a correspondingly shallow
/// autodiff graph). The sequential reference it replaces is kept as a test oracle
/// (`quat_cumprod_sequential` in the tests module) and asserted equal on values
/// **and** gradients.
pub fn quat_cumprod(q_bshj4: Tensor<5>, init: Option<Tensor<4>>) -> (Tensor<5>, Tensor<4>) {
    let [batch, sequence, nheads, blocks, _four] = q_bshj4.dims();
    let device = q_bshj4.device();

    // Pure prefix product Pₜ = qₜ ⊗ qₜ₋₁ ⊗ ⋯ ⊗ q₀ by Hillis–Steele doubling.
    // Invariant after each step with offset `d`: a[t] holds the product of the
    // window [t .. max(t-2d+1, 0)] (newest on the left). After ⌈log₂ sequence⌉
    // doublings the window covers [t .. 0], i.e. a[t] = Pₜ.
    let mut a = q_bshj4;
    let mut offset = 1usize;
    while offset < sequence {
        // shifted[t] = a[t-offset] for t ≥ offset, else the identity quaternion
        // (1,0,0,0) — so the first `offset` prefixes pass through unchanged
        // (a ⊗ identity = a).
        let ident = {
            let w = Tensor::ones([batch, offset, nheads, blocks, 1], &device);
            let xyz = Tensor::zeros([batch, offset, nheads, blocks, 3], &device);
            Tensor::cat(vec![w, xyz], 4)
        };
        let shifted = Tensor::cat(vec![ident, a.clone()], 1).narrow(1, 0, sequence);
        // Recent block (a) on the left, older block (shifted) on the right.
        a = quat_mul(a, shifted);
        offset *= 2;
    }

    // Fold the cross-chunk carry once: cumₜ = Pₜ ⊗ init. `init` (the previous
    // chunk's final cumulative rotation) is the oldest factor, hence on the
    // right; a missing carry is the identity and needs no multiply.
    let cum = match init {
        Some(init_bhj4) => {
            assert_eq!([batch, nheads, blocks, 4], init_bhj4.dims());
            quat_mul(a, init_bhj4.unsqueeze_dim::<5>(1)) // [batch, 1, nheads, J, 4] broadcasts over seq
        }
        None => a,
    };

    let final_carry = cum.clone().narrow(1, sequence - 1, 1).squeeze_dim::<4>(1); // [batch, nheads, J, 4]
    (cum, final_carry)
}

// ---------------------------------------------------------------------------
// Partial block rotation (rope_fraction support)
// ---------------------------------------------------------------------------

/// Apply a per-block quaternion rotation to the first `rope_width` entries of
/// the `state_rank` axis (a multiple of 4); the remainder passes through. The
/// quaternion analogue of [`apply_rope_partial`].
///
/// `q` has one quaternion per rotated block (`rope_width / 4` of them). `DB`
/// must equal `D + 1`.
pub fn rotate_blocks_partial<const D: usize, const DB: usize>(
    v: Tensor<D>,
    q: Tensor<DB>,
    rope_width: usize,
) -> Tensor<D> {
    let r = v.dims()[D - 1];
    debug_assert!(
        rope_width > 0,
        "a quaternion kind always turns at least one block; a real transition is RotationKind::Real1D"
    );
    if rope_width == r {
        rotate_state_rank_blocks::<D, DB>(v, q)
    } else {
        let head = v.clone().narrow(D - 1, 0, rope_width);
        let tail = v.narrow(D - 1, rope_width, r - rope_width);
        let head_rot = rotate_state_rank_blocks::<D, DB>(head, q);
        Tensor::cat(vec![head_rot, tail], D - 1)
    }
}

/// Two-sided counterpart of [`rotate_blocks_partial`]: rotates the first
/// `rope_width` entries of the `state_rank` axis by `v ↦ ql ⊗ v ⊗ qr`, passing
/// the remainder through.
pub fn rotate_blocks_two_sided_partial<const D: usize, const DB: usize>(
    v: Tensor<D>,
    ql: Tensor<DB>,
    qr: Tensor<DB>,
    rope_width: usize,
) -> Tensor<D> {
    let r = v.dims()[D - 1];
    debug_assert!(rope_width > 0, "see rotate_blocks_partial");
    if rope_width == r {
        rotate_state_rank_blocks_two_sided::<D, DB>(v, ql, qr)
    } else {
        let head = v.clone().narrow(D - 1, 0, rope_width);
        let tail = v.narrow(D - 1, rope_width, r - rope_width);
        let head_rot = rotate_state_rank_blocks_two_sided::<D, DB>(head, ql, qr);
        Tensor::cat(vec![head_rot, tail], D - 1)
    }
}

// ---------------------------------------------------------------------------
// Per-step rotation increments (the single definition every site derives from)
// ---------------------------------------------------------------------------

/// The abelian per-step angle increment `θ̂ₜ = Δₜ · range·π·tanh(ϑₜ)`, one per
/// head and rotation pair.
///
/// `DP1 = D + 1`: the head axis is inserted before the pair axis, so a
/// sequence-shaped call (`rot [b, s, a]`, `dt [b, s, h]`) yields `[b, s, h, a]`
/// and a single-token call (`rot [b, a]`, `dt [b, h]`) yields `[b, h, a]`.
pub fn angle_increment<const D: usize, const DP1: usize>(
    rot: Tensor<D>,
    dt: Tensor<D>,
    range: f64,
) -> Tensor<DP1> {
    assert_eq!(D + 1, DP1, "angle_increment inserts the head axis");
    let bounded = rot.tanh() * (range * std::f64::consts::PI);
    dt.unsqueeze_dim::<DP1>(D) * bounded.unsqueeze_dim::<DP1>(D - 1)
}

/// The quaternion per-step rotation vector `gₜ = Δₜ · range·π·tanh(‖r‖)·r̂`, one
/// per head and quaternion block (feed it to [`quat_from_scaled_axis`]).
///
/// The magnitude — not each component — is bounded, so the axis is exactly the
/// direction of the projection; see [`bound_rotation_vector`].
///
/// Unlike the abelian [`angle_increment`], the generators are projected **per
/// head**: `rot` carries `nheads · 3 · J` channels, so every head turns about
/// its own data-dependent axis rather than sharing one axis and differing only
/// in `Δ`. Sharing would make the heads' rotations a one-parameter family of
/// each other — the abelian path can afford that (its rotations commute, so a
/// per-head angle is the only freedom there is), but for a non-abelian
/// transition the axis *is* the expressive part: two heads turning about
/// different axes track different words, which is the whole point of having
/// more than one.
///
/// `DP1 = D + 1`, `DP2 = D + 2`: a sequence-shaped call (`rot [b, s, h·3·J]`,
/// `dt [b, s, h]`) yields `[b, s, h, J, 3]` and a single-token call
/// (`rot [b, h·3·J]`, `dt [b, h]`) yields `[b, h, J, 3]`.
pub fn generator_increment<const D: usize, const DP1: usize, const DP2: usize>(
    rot: Tensor<D>,
    dt: Tensor<D>,
    blocks: usize,
    range: f64,
) -> Tensor<DP2> {
    assert_eq!(D + 1, DP1, "generator_increment splits (3·J) into (J, 3)");
    assert_eq!(
        D + 2,
        DP2,
        "generator_increment also splits off the head axis"
    );
    let dims = rot.dims();
    let nheads = dt.dims()[D - 1];
    assert_eq!(
        dims[D - 1],
        nheads * 3 * blocks,
        "the rotation channels are three generators per (head, quaternion block)"
    );
    // [..., h·3·J] → [..., h, J, 3]
    let mut split = [0usize; DP2];
    split[..D - 1].copy_from_slice(&dims[..D - 1]);
    split[DP2 - 3] = nheads;
    split[DP2 - 2] = blocks;
    split[DP2 - 1] = 3;
    let bounded = bound_rotation_vector::<DP2>(rot.reshape(split), range * std::f64::consts::PI);
    // Δ is per head: broadcast it over (J, 3)
    bounded * dt.unsqueeze_dim::<DP1>(D).unsqueeze_dim::<DP2>(D + 1)
}

// ---------------------------------------------------------------------------
// Forward / step rotation of B and C (shared by both SSD pathways)
// ---------------------------------------------------------------------------

/// Rotate `B`/`C` for a **full sequence** by the data-dependent transition
/// rotation, returning the rotated projections and the new cumulative
/// [`RotationState`] to store in the cache.
///
/// Branches on [`RotationKind`]:
/// - [`Real1D`](RotationKind::Real1D): nothing happens — `B`/`C` pass through
///   and the (empty) accumulator is handed back. `rot` is `None` there, since
///   the block projects no rotation channels at all.
/// - [`Complex2D`](RotationKind::Complex2D): the abelian RoPE — cumulative
///   angle `cumsum` continued from `prev`, then [`apply_rope_partial`]. Exactly
///   the original Mamba-3 behaviour.
/// - [`Quaternion4D`](RotationKind::Quaternion4D): per-step unit quaternion
///   [`quat_from_scaled_axis`] (the in-projection generators scaled per-head by
///   `Δ`), composed by [`quat_cumprod`] continuing the cached quaternion, then
///   applied to `B`/`C` as `rotate(·, conj(Qₜ))` over the first `4·blocks`
///   state-rank entries.
///
/// # Shapes
/// - `rot_bsa` : `[batch, sequence, num_rotation_channels]` — the in-projection
///   rotation channels (angles for Complex2D, `3·blocks` quaternion generators
///   for Quaternion4D), `None` for Real1D, which projects none.
/// - `dt_bsh`  : `[batch, sequence, nheads]` (`Δ`).
/// - `b_bsmhr` / `c_bsmhr` : `[batch, sequence, mimo_rank, nheads, state_rank]`.
pub fn rotate_bc_forward(
    rot_bsa: Option<Tensor<3>>,
    dt_bsh: Tensor<3>,
    prev: RotationState,
    b_bsmhr: Tensor<5>,
    c_bsmhr: Tensor<5>,
    spec: RotationSpec,
) -> (Tensor<5>, Tensor<5>, RotationState) {
    let [batch, sequence, mimo_rank, nheads, _state_rank] = b_bsmhr.dims();
    let RotationSpec {
        kind,
        rope_dim,
        range,
    } = spec;
    // Only `Real1D` projects no rotation channels, and it never reads them.
    let rot = |r: Option<Tensor<3>>| r.expect("a rotating kind projects rotation channels");
    match kind {
        RotationKind::Real1D => (b_bsmhr, c_bsmhr, prev.expect_real()),
        RotationKind::Complex2D => {
            let rot_bsa = rot(rot_bsa);
            let prev_angle_bha = prev.angle();
            let num_rope_angles = prev_angle_bha.dims()[2];
            let raw_angles_bsha = angle_increment::<3, 4>(rot_bsa, dt_bsh, range);
            let cum_angles_bsha = prev_angle_bha.unsqueeze_dim::<4>(1) + raw_angles_bsha.cumsum(1);
            let cum_angles_bsmha = cum_angles_bsha.clone().unsqueeze_dim::<5>(2).expand([
                batch,
                sequence,
                mimo_rank,
                nheads,
                num_rope_angles,
            ]);
            let rotate_pairwise = mimo_rank == 1;
            let b = apply_rope_partial::<5>(
                b_bsmhr,
                cum_angles_bsmha.clone(),
                rope_dim,
                rotate_pairwise,
            );
            let c = apply_rope_partial::<5>(c_bsmhr, cum_angles_bsmha, rope_dim, rotate_pairwise);
            let last = wrap_angle(
                cum_angles_bsha
                    .narrow(1, sequence - 1, 1)
                    .squeeze_dim::<3>(1),
            );
            (b, c, RotationState::Angle(last))
        }
        RotationKind::Quaternion4D | RotationKind::Rotor4D => {
            // `stack` is the generator/scan block axis: `blocks` for
            // Quaternion4D, `2·blocks` for Rotor4D (left factors, then right).
            // Everything up to the *application* is factor-agnostic and runs
            // once over the whole stack.
            let (prev_q_bhk4, blocks) = prev.quat_stack(kind);
            let stack = prev_q_bhk4.dims()[2];
            // The rotated width is a whole number of quaternion blocks, and
            // the cache was sized for exactly that many.
            assert_eq!(rope_dim, blocks * 4, "cache/block rotation width mismatch");
            let rope_width = rope_dim;
            // Generators [b,s,h,stack,3]: the raw channels bounded to an angle
            // of at most `range·π` (see `generator_increment`) and scaled
            // per-head by Δ. The bound is load-bearing, not cosmetic: an
            // unbounded `g = rot·Δ` overflows f32 to `inf` for a large
            // in-projection activation, and `quat_from_scaled_axis`'s `cos(∞)`
            // then yields a forward NaN.
            let g_bshk3 = generator_increment::<3, 4, 5>(rot(rot_bsa), dt_bsh, stack, range);
            let q_step_bshk4 = quat_from_scaled_axis::<5>(g_bshk3);
            // Memory-efficient scan: a custom recompute backward (saves only the
            // leaf inputs) instead of retaining the scan's intermediates. Equal
            // to [`quat_cumprod`] on values and gradients (asserted in tests).
            let (cum_bshk4, final_bhk4) = crate::mamba3::quat_scan::quat_cumprod_recalculated(
                q_step_bshk4,
                Some(prev_q_bhk4),
            );
            // Renormalise the prefixes. A product of unit quaternions is a unit
            // quaternion in exact arithmetic, but the scan composes each prefix
            // out of `⌈log₂ L⌉` multiplies, so the norm drifts — negligibly in
            // f32, by ~1% over a long f16 sequence, and a non-unit rotation
            // rescales B/C instead of only turning them. `step` already
            // normalises every step; this is the chunkwise counterpart, and it
            // keeps the two numerically alike as well as exactly orthogonal.
            let cum_bshk4 = quat_normalize(cum_bshk4);
            let over_mimo = |q_bshj4: Tensor<5>| {
                q_bshj4
                    .unsqueeze_dim::<6>(2)
                    .expand([batch, sequence, mimo_rank, nheads, blocks, 4])
            };
            // B̄ = P⁻¹B: rotate by the inverse cumulative rotation, per block,
            // broadcast over the mimo_rank axis.
            let (b, c) = match kind {
                RotationKind::Rotor4D => {
                    // P(v) = Q v T̄  ⇒  P⁻¹(v) = Q* v T (conjugate on the left
                    // factor only).
                    let (left, right) = split_rotor(cum_bshk4);
                    let ql = over_mimo(quat_conj(left));
                    let qr = over_mimo(right);
                    (
                        rotate_blocks_two_sided_partial::<5, 6>(
                            b_bsmhr,
                            ql.clone(),
                            qr.clone(),
                            rope_width,
                        ),
                        rotate_blocks_two_sided_partial::<5, 6>(c_bsmhr, ql, qr, rope_width),
                    )
                }
                _ => {
                    let conj_bsmhj4 = over_mimo(quat_conj(cum_bshk4));
                    (
                        rotate_blocks_partial::<5, 6>(b_bsmhr, conj_bsmhj4.clone(), rope_width),
                        rotate_blocks_partial::<5, 6>(c_bsmhr, conj_bsmhj4, rope_width),
                    )
                }
            };
            let final_bhk4 = quat_normalize(final_bhk4);
            let state = match kind {
                RotationKind::Rotor4D => RotationState::Rotor(final_bhk4),
                _ => RotationState::Quaternion(final_bhk4),
            };
            (b, c, state)
        }
    }
}

/// Single-token counterpart of [`rotate_bc_forward`] for the recurrent `step`.
///
/// # Shapes
/// - `rot_ba`  : `[batch, num_rotation_channels]`, `None` for Real1D.
/// - `dt_bh`   : `[batch, nheads]`.
/// - `b_bmhr` / `c_bmhr` : `[batch, mimo_rank, nheads, state_rank]`.
pub fn rotate_bc_step(
    rot_ba: Option<Tensor<2>>,
    dt_bh: Tensor<2>,
    prev: RotationState,
    b_bmhr: Tensor<4>,
    c_bmhr: Tensor<4>,
    spec: RotationSpec,
) -> (Tensor<4>, Tensor<4>, RotationState) {
    let [batch, mimo_rank, nheads, _state_rank] = b_bmhr.dims();
    let RotationSpec {
        kind,
        rope_dim,
        range,
    } = spec;
    let rot = |r: Option<Tensor<2>>| r.expect("a rotating kind projects rotation channels");
    match kind {
        RotationKind::Real1D => (b_bmhr, c_bmhr, prev.expect_real()),
        RotationKind::Complex2D => {
            let rot_ba = rot(rot_ba);
            let prev_angle_bha = prev.angle();
            let num_rope_angles = prev_angle_bha.dims()[2];
            let raw_angle_bha = angle_increment::<2, 3>(rot_ba, dt_bh, range);
            let new_cum_angle_bha = wrap_angle(prev_angle_bha + raw_angle_bha);
            let new_cum_angle_bmha = new_cum_angle_bha.clone().unsqueeze_dim::<4>(1).expand([
                batch,
                mimo_rank,
                nheads,
                num_rope_angles,
            ]);
            let rotate_pairwise = mimo_rank == 1;
            let b = apply_rope_partial::<4>(
                b_bmhr,
                new_cum_angle_bmha.clone(),
                rope_dim,
                rotate_pairwise,
            );
            let c = apply_rope_partial::<4>(c_bmhr, new_cum_angle_bmha, rope_dim, rotate_pairwise);
            (b, c, RotationState::Angle(new_cum_angle_bha))
        }
        RotationKind::Quaternion4D | RotationKind::Rotor4D => {
            let (prev_q_bhk4, blocks) = prev.quat_stack(kind);
            let stack = prev_q_bhk4.dims()[2];
            assert_eq!(rope_dim, blocks * 4, "cache/block rotation width mismatch");
            let rope_width = rope_dim;
            // The same bounded generator as `rotate_bc_forward` (see the note there).
            let g_bhk3 = generator_increment::<2, 3, 4>(rot(rot_ba), dt_bh, stack, range);
            let q_step_bhk4 = quat_from_scaled_axis::<4>(g_bhk3);
            // Single step: Qₜ = qₜ ⊗ Qₜ₋₁ — and, stacked alongside it for
            // Rotor4D, Tₜ = pₜ ⊗ Tₜ₋₁ (the same left fold, see [`RotationState::Rotor`]).
            let new_q_bhk4 = quat_normalize(quat_mul(q_step_bhk4, prev_q_bhk4));
            let over_mimo = |q_bhj4: Tensor<4>| {
                q_bhj4
                    .unsqueeze_dim::<5>(1)
                    .expand([batch, mimo_rank, nheads, blocks, 4])
            };
            match kind {
                RotationKind::Rotor4D => {
                    let (left, right) = split_rotor(new_q_bhk4.clone());
                    let ql = over_mimo(quat_conj(left));
                    let qr = over_mimo(right);
                    let b = rotate_blocks_two_sided_partial::<4, 5>(
                        b_bmhr,
                        ql.clone(),
                        qr.clone(),
                        rope_width,
                    );
                    let c = rotate_blocks_two_sided_partial::<4, 5>(c_bmhr, ql, qr, rope_width);
                    (b, c, RotationState::Rotor(new_q_bhk4))
                }
                _ => {
                    let conj_bmhj4 = over_mimo(quat_conj(new_q_bhk4.clone()));
                    let b = rotate_blocks_partial::<4, 5>(b_bmhr, conj_bmhj4.clone(), rope_width);
                    let c = rotate_blocks_partial::<4, 5>(c_bmhr, conj_bmhj4, rope_width);
                    (b, c, RotationState::Quaternion(new_q_bhk4))
                }
            }
        }
    }
}

#[cfg(all(test, feature = "_dev-test"))]
mod tests;