Mixers¶

The five sequence mixers and the two container modules. Signatures and parameter contracts are generated from the source docstrings; for why each mixer behaves as it does, see the explanations, and for which to use, Choose a mixer.

mingru.MinGRU ¶

MinGRU(input_size: int, hidden_size: int, bias: bool = True, learnable_h0: bool = False, decay: Decay = None, decay_rate: float = 1.0, log1p_delta: bool = False)

Bases: DecayMixin, Module

minGRU layer: log-space parallel-scan training, recurrent inference.

Parameters:

Name	Type	Description	Default
`input_size`	`int`	Dimensionality of the inputs `x_t`.	required
`hidden_size`	`int`	Dimensionality of the hidden states `h_t`.	required
`bias`	`bool`	Whether the two linear maps carry bias terms.	`True`
`learnable_h0`	`bool`	If True, the module owns a learned initial state, stored as an unconstrained pre-activation and mapped through `g` internally (so positivity is by construction). Used whenever `forward()` or `step()` receive no explicit state. Zero-init gives `g(0) = 0.5`, matching the fixed default.	`False`
`decay`	`(fixed, learnable, None)`	Exponential time decay of the carried state, driven by `delta_t`: `gamma = exp(-lambda * f(delta_t))` multiplies the log-space transition coefficient (`log(1 - z)`) additively in log-space, i.e. `log(gamma * (1 - z)) = log(1 - z) - lambda * f(delta_t)` — no new exp/log round-trip. `None` disables decay: `delta_t` must then be omitted, and behavior is bit-identical to the module without this feature. `"fixed"`: `lambda = decay_rate`, a scalar buffer uniform across channels. `"learnable"`: `lambda = softplus(rho)`, one `rho` per hidden channel, initialized so `lambda == decay_rate` at construction. `delta_t = 0` gives `gamma = 1` exactly (no decay), with no special-casing of `t = 0`: callers pass `delta_t = 0` at a true sequence start, and chunked calls pass the real gap at the chunk boundary to match a full forward exactly.	`"fixed"`
`decay_rate`	`float`	Fixed decay rate, or the learnable rate's init target.	`1.0`
`log1p_delta`	`bool`	If True, `delta_t` is passed through `log1p` before scaling by `lambda` (compresses large gaps). See `_normalize_delta_t`. Note: negative/NaN/inf `delta_t` entries are always sanitized to finite non-negative values; the courtesy warning about them fires on CPU only (on CUDA they are fixed silently — no host sync, same rationale as the `h_0` async validation).	`False`

Raises:

Type	Description
`ValueError`	If `decay` is not one of `None`, `"fixed"`, `"learnable"` (at construction); or if `decay` is enabled without `delta_t`, or `delta_t` is given without decay enabled (at call time, in `forward`/`step`).

Notes

Shapes: forward maps x (B, T, input_size) with optional h_0 (B, 1, hidden_size) to (B, T, hidden_size); step maps x_t (B, input_size) with optional h_prev (B, hidden_size) to (B, hidden_size). delta_t follows the same optionality: (B, T) or (B, T, 1) for forward, (B,) or (B, 1) for step. Invalid delta_t entries (negative/NaN/inf) are always sanitized to finite, non-negative values; the courtesy warning about them is CPU-only (silent on CUDA, avoiding a host sync) — the sanitizing clamp itself always applies, on every device.

References

Feng et al., "Were RNNs All We Needed?", arXiv:2410.01201. Heinsen (2023) -- the log-space parallel scan used by forward.

Examples:

>>> import torch
>>> from mingru import MinGRU
>>> layer = MinGRU(input_size=32, hidden_size=64)
>>> x = torch.randn(4, 128, 32)
>>> h = layer(x)                       # parallel training-mode forward
>>> tuple(h.shape)
(4, 128, 64)
>>> h_t = layer.step(x[:, 0])          # one O(1)-memory streaming step
>>> tuple(h_t.shape)
(4, 64)

forward ¶

forward(x: Tensor, h_0: Tensor | None = None, delta_t: Tensor | None = None) -> torch.Tensor

Parallel (training-mode) forward over a full sequence.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input sequence, shape `(B, T, input_size)`.	required
`h_0`	`Tensor`	Initial hidden state, shape `(B, 1, hidden_size)`, as a REAL state — i.e. an output of `forward()`/`step()`, non-negative (exact zeros from fp16/bf16 underflow are clamped). Do NOT pass a pre-activation; `g` is not applied here (this differs from the paper's reference code, which treats `h_0` as a pre-activation). Defaults to `g(0) = 0.5`, matching `step()`'s default. For chunked/TBPTT training, carry `h_0 = prev_out[:, -1:]` (detach as appropriate).	`None`
`delta_t`	`Tensor`	Time gaps preceding each event, shape `(B, T)` or `(B, T, 1)`; required iff `decay` is enabled (see class docstring). `delta_t[:, t]` is the gap preceding event `t`; `delta_t = 0` means no decay at that step.	`None`

Returns:

Type	Description
`Tensor`	All hidden states `h_1..h_T`, shape `(B, T, hidden_size)`.

Raises:

Type	Description
`RuntimeError`	If `h_0` contains strictly negative entries (the signature of pre-activation misuse). Raised device-side via `torch._assert_async`; on CUDA the error surfaces at the next sync point rather than at this call.
`ValueError`	If `decay` is enabled without `delta_t`, or `delta_t` is given without decay enabled.

step ¶

step(x_t: Tensor, h_prev: Tensor | None = None, delta_t: Tensor | None = None) -> torch.Tensor

Single recurrent step for streaming/token-by-token inference.

Numerically equivalent to forward() one timestep at a time (g is applied to the candidate to match the log-space parameterization). h_prev is a real hidden state — the same convention as forward()'s h_0.

Parameters:

Name	Type	Description	Default
`x_t`	`Tensor`	Input at the current timestep, shape `(B, input_size)`.	required
`h_prev`	`Tensor`	Previous hidden state, shape `(B, hidden_size)` — an output of `step()` or `forward()`. Defaults to `g(0) = 0.5` (or the learned initial state if `learnable_h0`), matching `forward()`'s default.	`None`
`delta_t`	`Tensor`	Time gap preceding this event, shape `(B,)` or `(B, 1)`; required iff `decay` is enabled.	`None`

Returns:

Type	Description
`Tensor`	New hidden state `h_t`, shape `(B, hidden_size)`.

Raises:

Type	Description
`ValueError`	If `decay` is enabled without `delta_t`, or `delta_t` is given without decay enabled.

mingru.SignedMinGRU ¶

SignedMinGRU(input_size: int, hidden_size: int, bias: bool = True, learnable_h0: bool = False, coupled: bool = False, decay: Decay = None, decay_rate: float = 1.0, log1p_delta: bool = False)

Bases: DecayMixin, Module

minGRU variant with signed diagonal transitions (linear-space scan).

The recurrence is h_t = a_t * h_{t-1} + z_t * h~_t with a_t ranging over (-1, 1) instead of (0, 1). Motivated by Merrill, Petty & Sabharwal (2024): still diagonal, hence commutative and TC0, but negative eigenvalues restore per-layer parity / sign-alternation dynamics (the mechanism identified by Grazzi et al., 2025).

Two parameterizations of a_t are available:

coupled=False (default): a_t = tanh(Linear_s(x_t)), the eigenvalue decoupled from the update gate z_t. This is the experimentally superior form (parity accuracy @1024: 0.61 -> 0.996, 6-seed mean, current-env) and is now the default.
coupled=True: a_t = (1 - z_t) * tanh(Linear_s(x_t)), the original parameterization (eigenvalue coupled to the update gate, as in the vanilla minGRU's (1 - z_t) retention coefficient). Bit-exact reproduction of the pre-promotion class: identical parameter shapes and construction order (linear_z, linear_h, linear_s), so identical seeds give identical weights.

When the sign head saturates positive, the coupled form reduces to the vanilla (Appendix A) minGRU's retention dynamics (a = 1 - z); the decoupled form instead saturates to a = 1, a perfect integrator — reaching the interval boundary is exactly what the coupling prevents, and why the decoupled form length-generalizes.

States are unconstrained reals: no g, no positivity checks, no underflow handling. Same forward/step API and shapes as MinGRU; h_0 is any real state.

Parameters:

Name	Type	Description	Default
`input_size`	`int`	Dimensionality of the inputs `x_t`.	required
`hidden_size`	`int`	Dimensionality of the hidden states `h_t`.	required
`bias`	`bool`	Whether the three linear maps carry bias terms.	`True`
`learnable_h0`	`bool`	If True, the module owns a learned initial state (an unconstrained parameter used directly — no `g`). Zero-init matches the fixed default `h_0 = 0`.	`False`
`coupled`	`bool`	If True, use the legacy coupled eigenvalue `a_t = (1 - z_t) * tanh(Linear_s(x_t))`. If False (default), use the decoupled eigenvalue `a_t = tanh(Linear_s(x_t))`.	`False`
`decay`	`(fixed, learnable, None)`	Exponential time decay of the carried state: `a_decayed = gamma * a` with `gamma = exp(-lambda * f(delta_t))` — magnitude decays, sign is preserved, applied identically for both `coupled` values. `None` disables decay: `delta_t` must then be omitted, and behavior (including the `coupled=True` bit-exact reproduction guarantee) is bit-identical to the module without this feature. `"fixed"`: `lambda = decay_rate`, a scalar buffer. `"learnable"`: `lambda = softplus(rho)`, one `rho` per hidden channel, initialized so `lambda == decay_rate` at construction. `delta_t = 0` gives `gamma = 1` exactly, with no special-casing of `t = 0`.	`"fixed"`
`decay_rate`	`float`	Fixed decay rate, or the learnable rate's init target.	`1.0`
`log1p_delta`	`bool`	If True, `delta_t` is passed through `log1p` before scaling by `lambda`. See `_normalize_delta_t`. Note: negative/NaN/inf `delta_t` entries are always sanitized to finite non-negative values; the courtesy warning about them fires on CPU only (on CUDA they are fixed silently — no host sync, same rationale as the `h_0` async validation).	`False`

Raises:

Type	Description
`ValueError`	If `decay` is not one of `None`, `"fixed"`, `"learnable"` (at construction); or if `decay` is enabled without `delta_t`, or `delta_t` is given without decay enabled (at call time, in `forward`/`step`).

Notes

Parameter count is 3 linear heads vs. MinGRU's 2 (the extra sign head) — account for this in parameter-matched comparisons. Invalid delta_t entries (negative/NaN/inf) are always sanitized to finite, non-negative values; the courtesy warning about them is CPU-only (silent on CUDA, avoiding a host sync) — the sanitizing clamp itself always applies, on every device.

References

Merrill, Petty & Sabharwal (2024) -- expressivity gains from signed (negative-eigenvalue) diagonal linear recurrences. Grazzi et al. (2025) -- sign-alternation / parity dynamics enabled by negative transition eigenvalues.

Examples:

>>> import torch
>>> from mingru import SignedMinGRU
>>> layer = SignedMinGRU(input_size=32, hidden_size=64)  # decoupled default
>>> x = torch.randn(4, 128, 32)
>>> tuple(layer(x).shape)
(4, 128, 64)

forward ¶

forward(x: Tensor, h_0: Tensor | None = None, delta_t: Tensor | None = None) -> torch.Tensor

Parallel forward over a full sequence.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input sequence, shape `(B, T, input_size)`.	required
`h_0`	`Tensor`	Initial hidden state, shape `(B, 1, hidden_size)`. Any real values. Defaults to zeros (or the learned initial state if `learnable_h0`).	`None`
`delta_t`	`Tensor`	Time gaps preceding each event, shape `(B, T)` or `(B, T, 1)`; required iff `decay` is enabled.	`None`

Returns:

Type	Description
`Tensor`	All hidden states `h_1..h_T`, shape `(B, T, hidden_size)`.

Raises:

Type	Description
`ValueError`	If `decay` is enabled without `delta_t`, or `delta_t` is given without decay enabled.

step ¶

step(x_t: Tensor, h_prev: Tensor | None = None, delta_t: Tensor | None = None) -> torch.Tensor

Single recurrent step; same real-state convention as forward().

Computed from the same _coeffs helper forward() uses, so decayed and non-decayed dynamics cannot drift between the two call paths.

Parameters:

Name	Type	Description	Default
`x_t`	`Tensor`	Input at the current timestep, shape `(B, input_size)`.	required
`h_prev`	`Tensor`	Previous hidden state, shape `(B, hidden_size)`. Defaults to zeros (or the learned initial state).	`None`
`delta_t`	`Tensor`	Time gap preceding this event, shape `(B,)` or `(B, 1)`; required iff `decay` is enabled.	`None`

Returns:

Type	Description
`Tensor`	New hidden state, shape `(B, hidden_size)`.

Raises:

Type	Description
`ValueError`	If `decay` is enabled without `delta_t`, or `delta_t` is given without decay enabled.

mingru.RotationMinGRU ¶

RotationMinGRU(input_size: int, hidden_size: int, bias: bool = True, snap: tuple[int, ...] | None = (2, 3, 4, 6), decay: Decay = None, decay_rate: float = 1.0, log1p_delta: bool = False)

Bases: DecayMixin, Module

minGRU variant with 2x2 block rotation transitions (non-diagonal).

State is n = hidden_size / 2 independent planar (2D) blocks. Per block, the recurrence is a full 2x2 affine map instead of a scalar one:

M_t = R(theta_t) @ diag(1, tanh(u_t))
h_t = M_t @ h_{t-1} + b_t,     b_t = z_t * Linear_h(x_t)

(h_{t-1}, h_t, b_t viewed as 2-vectors per block). Unlike MinGRU and SignedMinGRU, per-block transitions do not commute (2x2 rotation/reflection matrices form a non-abelian group under composition), so this mixer — with a non-commutative parallel scan, matrix_scan — can represent state-tracking automata over non-abelian groups that a diagonal (commutative) scan provably cannot. D3 (isomorphic to S3, the smallest non-abelian group) embeds in O(2), so one layer of this mixer can represent the S3 running product exactly; see experiments/SUMMARY.md for the mechanism verification (per-block matrices extracted from a trained model satisfy the D3 composition table to ~1e-4).

Angle snapping (snap): with snap set, theta_t is quantized per block to an exact multiple of 2*pi/K via a straight-through estimator — forward uses the snapped angle, gradient passes through the pre-snap "soft" angle unchanged. K is cycled across blocks from the snap tuple (block j uses snap[j % len(snap)]). This manufactures attractors at exact group elements -- exactly for the snapped angle; the scale channel tanh(u_t) reaches a group element (+-1) only as the u-head saturates, so full-map exactness is empirical -- the same way tanh's asymptote manufactures an attractor at eigenvalue -1 for SignedMinGRU: without snapping, plain rotation angles have no attractor and drift with sequence length (error compounds with T). The snap grid must contain the group being tracked: choose K values whose rotations (2*pi/K) generate, or coincide with, the target group's rotation subgroup, or the exact automaton is not representable on the grid at all (e.g. tracking Z/5 needs a multiple of 5 in snap). The default snap=(2, 3, 4, 6) was chosen for the D3/S3 task; other state-tracking targets need their own grid. snap=None gives continuous (unsnapped) rotations — a legitimate, documented ladder rung, but angles then drift under length generalization since there is no attractor at the task's true transition angle; use only when exact length generalization is not required.

Depth: single-layer use is the recorded baseline; stacks holding rotation blocks (with or without other mixer types) are validated at L=2 on the S3 probe under the best-val@128 protocol, deeper is untested. The straight-through discontinuity can compound across rotation layers, so multi-rotation stacks warn at construction — see MinGRUStack and the README's Rotation variant section.

Training protocol: the exact automaton is reachable but is NOT a stable attractor of standard training — runs wander in and out of it during optimization. The validated protocol is best-checkpoint selection by validation accuracy at a length LONGER than the training length (e.g. T=128 when training at T=64; not one of the eventual test lengths), evaluated over the full step budget instead of early-stopping, plus a retry-on-flag rule: a best validation score at that checkpoint length below 1.0 flags the run as failed. The flag is one-directional: a sub-1.0 score reliably marks a bad run, but a perfect checkpoint score does not guarantee exact length generalization (in the recorded evidence every seed passed the checkpoint yet most still decayed at the longest lengths). See experiments/SUMMARY.md for the full protocol, per-seed success rate, and mechanism verification.

Excludes refuted experiment-loop mechanisms: no full orthogonality constraint (ortho), no grid-attraction regularizer (reg), no post-hoc projection/ablation masks. All were tried and either hurt length generalization or were redundant with the best-val selection protocol above; see experiments/SUMMARY.md rounds 5 and 8.

h_0 is an unconditional learnable parameter (no learnable_h0 flag, unlike the module's other two mixers): h_0 = 0 has no orbit under the group action (a fixed point cannot demonstrate state tracking), and a state vector lying on a reflection axis collapses reflections onto rotations. A random nonzero learned vector avoids both failure modes.

Parameters:

Name	Type	Description	Default
`input_size`	`int`	Dimensionality of the inputs `x_t`.	required
`hidden_size`	`int`	Dimensionality of the hidden states `h_t`; must be even (`hidden_size = 2 * n_blocks`).	required
`bias`	`bool`	Whether the four linear maps carry bias terms.	`True`
`snap`	`tuple of int, or None`	Per-block angle-snap grid orders `K` (cycled across blocks); each block's angle snaps to multiples of `2*pi/K`. `None` disables snapping (continuous rotations; see above).	`(2, 3, 4, 6)`
`decay`	`(fixed, learnable, None)`	Exponential time decay of the carried state: `M_decayed = gamma * M` per block, with `gamma = exp(-lambda * f(delta_t))`. Scalar decay commutes with the rotation/ reflection group action, so the composed transition's ANGLE is unaffected (stays exactly on the snap grid) — only amplitude fades. `None` disables decay: `delta_t` must then be omitted, and behavior is bit-identical to the module without this feature. `"fixed"`: `lambda = decay_rate`, a scalar buffer. `"learnable"`: `lambda = softplus(rho)`, one `rho` per block (`n_blocks` channels), initialized so `lambda == decay_rate` at construction. `delta_t = 0` gives `gamma = 1` exactly, with no special-casing of `t = 0`.	`"fixed"`
`decay_rate`	`float`	Fixed decay rate, or the learnable rate's init target.	`1.0`
`log1p_delta`	`bool`	If True, `delta_t` is passed through `log1p` before scaling by `lambda`. See `_normalize_delta_t`. Note: negative/NaN/inf `delta_t` entries are always sanitized to finite non-negative values; the courtesy warning about them fires on CPU only (on CUDA they are fixed silently — no host sync, same rationale as the `h_0` async validation).	`False`

Raises:

Type	Description
`ValueError`	If `hidden_size` is odd, or `decay` is not one of `None`, `"fixed"`, `"learnable"` (at construction); or if `decay` is enabled without `delta_t`, or `delta_t` is given without decay enabled (at call time, in `forward`/`step`).

Notes

Parameter count is 4 linear heads (z, h, theta, u) vs. SignedMinGRU's 3 — account for this in parameter-matched comparisons. Same forward/step shapes as the module's other mixers: forward maps x (B, T, input_size) with optional h_0 (B, 1, hidden_size) to (B, T, hidden_size); step maps x_t (B, input_size) with optional h_prev (B, hidden_size) to (B, hidden_size). delta_t follows the same optionality: (B, T) or (B, T, 1) for forward, (B,) or (B, 1) for step. Invalid delta_t entries (negative/NaN/inf) are always sanitized to finite, non-negative values; the courtesy warning about them is CPU-only (silent on CUDA, avoiding a host sync) — the sanitizing clamp itself always applies, on every device.

References

Merrill, Petty & Sabharwal (2024) and Grazzi et al. (2025) -- state tracking over non-abelian groups needs non-commutative (matrix) transitions, which a diagonal scan cannot express.

Examples:

>>> import torch
>>> from mingru import RotationMinGRU
>>> layer = RotationMinGRU(input_size=32, hidden_size=64)  # snap=(2,3,4,6)
>>> x = torch.randn(4, 128, 32)
>>> tuple(layer(x).shape)
(4, 128, 64)

forward ¶

forward(x: Tensor, h_0: Tensor | None = None, delta_t: Tensor | None = None) -> torch.Tensor

Parallel forward over a full sequence.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input sequence, shape `(B, T, input_size)`.	required
`h_0`	`Tensor`	Initial hidden state, shape `(B, 1, hidden_size)`. Any real values (reshaped into `n_blocks` 2-vectors internally). Defaults to the learned initial state.	`None`
`delta_t`	`Tensor`	Time gaps preceding each event, shape `(B, T)` or `(B, T, 1)`; required iff `decay` is enabled.	`None`

Returns:

Type	Description
`Tensor`	All hidden states `h_1..h_T`, shape `(B, T, hidden_size)`.

Raises:

Type	Description
`ValueError`	If `decay` is enabled without `delta_t`, or `delta_t` is given without decay enabled.

step ¶

step(x_t: Tensor, h_prev: Tensor | None = None, delta_t: Tensor | None = None) -> torch.Tensor

Single recurrent step; same real-state convention as forward().

Computed from the same _coeffs helper forward() uses (applied per-step instead of over the full sequence), so the two paths cannot drift apart — mirrors how SignedMinGRU shares _coeffs between its forward/step.

Parameters:

Name	Type	Description	Default
`x_t`	`Tensor`	Input at the current timestep, shape `(B, input_size)`.	required
`h_prev`	`Tensor`	Previous hidden state, shape `(B, hidden_size)`. Defaults to the learned initial state.	`None`
`delta_t`	`Tensor`	Time gap preceding this event, shape `(B,)` or `(B, 1)`; required iff `decay` is enabled.	`None`

Returns:

Type	Description
`Tensor`	New hidden state, shape `(B, hidden_size)`.

Raises:

Type	Description
`ValueError`	If `decay` is enabled without `delta_t`, or `delta_t` is given without decay enabled.

mingru.GivensMinGRU ¶

GivensMinGRU(input_size: int, hidden_size: int, bias: bool = True, block_size: int = 8, rounds: int = 3, decay: Decay = None, decay_rate: float = 1.0, log1p_delta: bool = False)

Bases: DecayMixin, Module

minGRU variant with k-dim block-rotation transitions (non-diagonal).

The k > 2 generalization of RotationMinGRU's snap=None regime. State is n_blocks = hidden_size / block_size independent blocks of k = block_size dims (64 state elements total at the repo's d_model = 512: the same per-token state as every other promoted mixer). Per block and token the transition is a product of rounds layers of Givens rotations on fixed brick-wall planes,

M_t = G_{rounds-1} @ ... @ G_1 @ G_0,
h_t = M_t @ h_{t-1} + b_t,     b_t = sigmoid(z_t) * Linear_h(x_t)

(h_{t-1}, h_t, b_t viewed as k-vectors per block). Each round r rotates a set of disjoint coordinate planes by input-dependent angles from one linear head: even rounds pair planes (0,1),(2,3),...; odd rounds the staggered (1,2),(3,4),...,(k-1,0). Every G_r is orthogonal with determinant +1 by construction, so M_t is exactly special-orthogonal — continuous, with no straight-through snap. The brick-wall plane layout is the standard rectangular mesh from the orthogonal/unitary-RNN literature (EUNN: Jing et al. 2017, arXiv:1612.05231; mesh design: Clements et al. 2016) built from Givens plane rotations (Golub & Van Loan, Matrix Computations); the departure here is input-dependent angles per token.

Where RotationMinGRU manufactures attractors at exact group elements by snapping 2x2 angles, GivensMinGRU deliberately does not: each per-token map lives on a rounds * block_size / 2-angle submanifold of SO(k) — 12 of SO(8)'s 28 dimensions at the defaults; products of enough Givens rotations (about k - 1 brick-wall rounds) generate all of SO(k), so three rounds is a deliberate budget rather than full per-token coverage — giving richer non-abelian per-token maps at matched state capacity, at the cost of having no attractor at any particular transition — angles drift under length generalization exactly as RotationMinGRU(snap=None)'s do. Like the other matrix-transition mixers it runs on a non-commutative parallel scan (matrix_affine_scan, the k-dim generalization of matrix_scan), so it can represent non-abelian state tracking a diagonal (commutative) scan provably cannot.

Depth: unlike RotationMinGRU, stacks holding multiple Givens blocks construct WITHOUT a warning — the multi-rotation UserWarning names straight-through snap compounding, which the continuous (unsnapped) Givens transition does not have (see MinGRUStack).

h_0 is an unconditional learnable parameter (no learnable_h0 flag, same convention and rationale as RotationMinGRU): h_0 = 0 has no orbit under the group action, so a zero initial state cannot demonstrate state tracking; a random nonzero learned vector avoids that failure mode.

Parameters:

Name	Type	Description	Default
`input_size`	`int`	Dimensionality of the inputs `x_t`.	required
`hidden_size`	`int`	Dimensionality of the hidden states `h_t`; must be an integer multiple of `block_size` (`hidden_size = n_blocks * block_size`).	required
`bias`	`bool`	Whether the three linear maps carry bias terms.	`True`
`block_size`	`int`	Per-block state dimension `k`; must be even and divide `hidden_size`. The default keeps the standard 64-element per-token state at the repo's `d_model`.	`8`
`rounds`	`int`	Number of brick-wall Givens layers composed per transition; must be at least 1. Full `SO(k)` coverage would take `k * (k - 1) / 2` rotations (about `k - 1` brick-wall rounds); the default three rounds span a 12-angle submanifold of `SO(8)` and are the evidence-validated budget.	`3`
`decay`	`(fixed, learnable, None)`	Exponential time decay of the carried state: `M_decayed = gamma * M` per block, with `gamma = exp(-lambda * f(delta_t))`. A scalar `gamma` commutes with the orthogonal block action, so the composed transition's rotation is unaffected — only amplitude fades — identical decay semantics to `RotationMinGRU`. `None` disables decay: `delta_t` must then be omitted, and the forward computation is bit-identical to the module without this feature (no `gamma = 1` multiply on the disabled-decay path). `"fixed"`: `lambda = decay_rate`, a scalar buffer. `"learnable"`: `lambda = softplus(rho)`, one `rho` per block (`n_blocks` channels), initialized so `lambda == decay_rate` at construction. `delta_t = 0` gives `gamma = 1` exactly.	`"fixed"`
`decay_rate`	`float`	Fixed decay rate, or the learnable rate's init target.	`1.0`
`log1p_delta`	`bool`	If True, `delta_t` is passed through `log1p` before scaling by `lambda`. See `_normalize_delta_t`.	`False`

Raises:

Type	Description
`ValueError`	If `hidden_size` is not a multiple of `block_size`, or `block_size` is odd (at construction); or if `decay` is not one of `None`, `"fixed"`, `"learnable"` (at construction); or if `decay` is enabled without `delta_t`, or `delta_t` is given without decay enabled (at call time, in `forward`/`step`).

Notes

Same forward/step shapes as the module's other mixers: forward maps x (B, T, input_size) with optional h_0 (B, 1, hidden_size) to (B, T, hidden_size); step maps x_t (B, input_size) with optional h_prev (B, hidden_size) to (B, hidden_size). delta_t follows the same optionality and shapes as RotationMinGRU ((B, T) or (B, T, 1) for forward; (B,) or (B, 1) for step), with the identical decay/delta_t pairing contract (ValueError both directions, see _validate_delta_t_pairing) and CPU-only invalid-entry warning. Parameter count is 3 linear heads (theta, z, h): the angle head emits n_blocks * rounds * (block_size / 2) angles.

References

Jing et al. (2017), "Tunable Efficient Unitary Neural Networks (EUNN)", arXiv:1612.05231 -- the brick-wall mesh of plane rotations. Clements et al. (2016) -- the rectangular (brick-wall) mesh design. Golub & Van Loan, Matrix Computations -- Givens plane rotations.

Examples:

>>> import torch
>>> from mingru import GivensMinGRU
>>> layer = GivensMinGRU(input_size=32, hidden_size=64)  # block_size=8, rounds=3
>>> x = torch.randn(4, 128, 32)
>>> tuple(layer(x).shape)
(4, 128, 64)

forward ¶

forward(x: Tensor, h_0: Tensor | None = None, delta_t: Tensor | None = None) -> torch.Tensor

Parallel forward over a full sequence.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input sequence, shape `(B, T, input_size)`.	required
`h_0`	`Tensor`	Initial hidden state, shape `(B, 1, hidden_size)` (reshaped into `n_blocks` `k`-vectors internally). Defaults to the learned initial state.	`None`
`delta_t`	`Tensor`	Time gaps preceding each event, shape `(B, T)` or `(B, T, 1)`; required iff `decay` is enabled.	`None`

Returns:

Type	Description
`Tensor`	All hidden states `h_1..h_T`, shape `(B, T, hidden_size)`.

Raises:

Type	Description
`ValueError`	If `decay` is enabled without `delta_t`, or `delta_t` is given without decay enabled.

step ¶

step(x_t: Tensor, h_prev: Tensor | None = None, delta_t: Tensor | None = None) -> torch.Tensor

Single recurrent step; same real-state convention as forward().

Computed from the same _coeffs helper forward() uses (applied per-step instead of over the full sequence), so the two paths cannot drift apart — mirrors RotationMinGRU.

Parameters:

Name	Type	Description	Default
`x_t`	`Tensor`	Input at the current timestep, shape `(B, input_size)`.	required
`h_prev`	`Tensor`	Previous hidden state, shape `(B, hidden_size)`. Defaults to the learned initial state.	`None`
`delta_t`	`Tensor`	Time gap preceding this event, shape `(B,)` or `(B, 1)`; required iff `decay` is enabled.	`None`

Returns:

Type	Description
`Tensor`	New hidden state, shape `(B, hidden_size)`.

Raises:

Type	Description
`ValueError`	If `decay` is enabled without `delta_t`, or `delta_t` is given without decay enabled.

mingru.DeltaMinGRU ¶

DeltaMinGRU(input_size: int, hidden_size: int, bias: bool = True, n_heads: int = 4, nh: int = 1, d_k: int | None = None, d_v: int | None = None, chunk_size: int = 64, decay: Decay = None, decay_rate: float = 1.0, log1p_delta: bool = False)

Bases: Module

DeltaNet/DeltaProduct generalized-Householder delta-rule mixer.

Departs from the other four mixers' diagonal/rotation-scan mechanism entirely: state is a per-head d_k x d_v associative-memory matrix H (not a vector), updated by nh generalized-Householder rank-1 corrections per token rather than a linear/affine scan. Reimplements Yang et al.'s DeltaNet (nh=1) and Siems et al.'s DeltaProduct generalization (nh>1) as a package citizen, mirroring the parameter layout of the lab's DeltaNetMixer (experiments/variants.py, not imported here) so state dicts are relatable across lab and package.

Per head, per token t, in order for micro-step j = 1..nh:

k_{t,j} = normalize(Linear_k^(j)(x_t))              # L2, last dim
v_{t,j} = Linear_v^(j)(x_t)
beta_{t,j} = 2 * sigmoid(Linear_beta^(j)(x_t))        # in (0, 2)
H <- (I - beta_{t,j} k_{t,j} k_{t,j}^T) H + beta_{t,j} k_{t,j} v_{t,j}^T

beta = 2 with a unit k makes the update an exact Householder reflection (I - 2 k k^T is orthogonal, determinant -1); (0, 2) interpolates between the identity (beta=0, no update) and a full reflection. Readout after the token's final micro-step is y_t = H_t^T q_t per head (q_t = Linear_q(x_t)), heads concatenated and projected by out_proj to hidden_size. H is zero-initialized (empty associative memory) unless h_0 is supplied -- unlike RotationMinGRU/GivensMinGRU, there is no intrinsic learned initial state: H = 0 is the delta rule's natural origin, not a degenerate point the way it is for a group action.

step implements the sequential per-token recurrence above directly (correct, O(T * nh) per-sequence cost) and is forward's correctness oracle. forward instead runs the chunked-WY parallel form (the standard efficient algorithm for this recurrence -- Yang et al. 2024, arXiv:2406.06484): the nh * T ordered rank-1 updates are processed chunk_size tokens at a time, each chunk's cumulative transition carried via one unit-lower-triangular solve (the UT transform) instead of a per-token loop. Both forms are exactly the same function (semantically identical within floating-point tolerance, chunk_size has no effect on the result) and neither ever materializes a per-token (or per-chunk) d_k x d_k transition matrix: each micro-step is a rank-1 correction, and the chunked form's UT transform operates on nh * chunk_size-sized matrices built from k/v/beta, never on an explicit d_k x d_k matrix. The fp32 agreement is chunk-size-bounded at large T, though: at T = 1024 with the default chunk_size = 64 the sequential/chunked deviation is ~3e-6, but a single-chunk run (chunk_size >= T) at the same T accumulates to ~1.2e-5 through the one (nh * C)-deep forward substitution the UT transform solves in a single shot rather than T separate nh-deep ones -- still inside the repo's atol=1e-5 forward tolerance, but close enough to it that the default (not an oversized) chunk size is recommended for long sequences.

Time decay is not supported: the constructor rejects any decay is not None (decay/decay_rate/log1p_delta are accepted only for signature uniformity with the other mixers -- folding a per-token decay into the delta-rule composition is unimplemented), and both forward/step reject any delta_t unconditionally.

Parameters:

Name	Type	Description	Default
`input_size`	`int`	Dimensionality of the inputs `x_t`.	required
`hidden_size`	`int`	Output dimensionality (`out_proj` maps `n_heads * d_v -> hidden_size`).	required
`bias`	`bool`	Whether every linear head (`q`, each micro-step's `k`/`v`/ `beta`, and `out_proj`) carries a bias term.	`True`
`n_heads`	`int`	Number of parallel associative-memory heads.	`4`
`nh`	`int`	Number of generalized-Householder micro-steps per token (`nh=1` is DeltaNet; `nh>1` is DeltaProduct).	`1`
`d_k`	`int`	Per-head key/query dimension. Defaults to `hidden_size // n_heads` (requires `hidden_size % n_heads == 0` in that case); pass explicitly to decouple per-head state size from `hidden_size`.	`None`
`d_v`	`int`	Per-head value dimension. Defaults to `d_k`.	`None`
`chunk_size`	`int`	Number of tokens processed per UT-transform chunk in the chunked-WY `forward` (performance-only knob: outputs and gradients are invariant to `chunk_size`, including `C=1` and `C >= T` and ragged final chunks). Validated at construction (must be positive).	`64`
`decay`	`(fixed, learnable, None)`	Must be `None` -- time decay is not implemented for the delta-rule transition. Accepted for signature uniformity with the other mixers only.	`"fixed"`
`decay_rate`	`float`	Unused (`decay` must be `None`); accepted for signature uniformity.	`1.0`
`log1p_delta`	`bool`	Unused (`decay` must be `None`); accepted for signature uniformity.	`False`

Raises:

Type	Description
`ValueError`	If `n_heads`, `nh`, or `chunk_size` is not positive; if `d_k` is not given and `hidden_size` is not divisible by `n_heads`; or if `decay` is not `None` (at construction). If `delta_t` is given (`forward`/`step`, decay can never be enabled) or `h_0`/`h_prev` has the wrong shape (at call time).

Notes

carries_matrix_state = True (class attribute): the state forward(..., return_state=True) returns and step carries is the flattened H matrix, not a readout vector -- y_t != h_t, unlike the other four mixers, where the scanned state doubles as the output. forward maps x (B, T, input_size) with optional h_0 (B, 1, n_heads*d_k*d_v) (flattened in (n_heads, d_k, d_v) row-major order) to y (B, T, hidden_size), or (y, h_T) when return_state=True. step maps x_t (B, input_size) with optional h_prev (B, n_heads*d_k*d_v) to (y_t, h_t), each following the same flattening. Chaining forward(x2, h_0=s) after (y1, s) = forward(x1, return_state=True) equals forward(cat(x1, x2)) on both outputs and gradients.

On CUDA, torch.compile is the recommended execution path (it recovers most of the chunked-WY forward's fusion headroom with no engineering). A hand-written Triton kernel for the forward also exists: under MINGRU_SCAN=auto it engages only in a narrow, measured win region (long sequences with narrow head dims), where it is faster than eager and lighter on memory; MINGRU_SCAN=triton forces the full kernel envelope; MINGRU_SCAN=eager disables it. The kernel does not reach torch.compile-class speed and is not a substitute for it -- see the docs' "Choose a mixer" guide (GPU execution path) for the full picture.

References

Yang et al., "Parallelizing Linear Transformers with the Delta Rule over Sequence Length" (DeltaNet), arXiv:2406.06484. Siems et al., "DeltaProduct: Increasing the Expressivity of DeltaNet Through Products of Householders" (DeltaProduct).

Examples:

>>> import torch
>>> from mingru import DeltaMinGRU
>>> layer = DeltaMinGRU(input_size=32, hidden_size=64, n_heads=4)
>>> x = torch.randn(4, 128, 32)
>>> tuple(layer(x).shape)
(4, 128, 64)

forward ¶

forward(x: Tensor, h_0: Tensor | None = None, delta_t: Tensor | None = None, return_state: bool = False) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]

Chunked-WY forward over a full sequence.

Computes exactly the function step computes token-by-token (chunk_size is a performance-only knob -- see _forward_chunked), via the chunked-WY UT transform instead of a per-token Python loop.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input sequence, shape `(B, T, input_size)`.	required
`h_0`	`Tensor`	Initial per-head associative-memory state, flattened shape `(B, 1, n_headsd_kd_v)` in `(n_heads, d_k, d_v)` row-major order. Defaults to zero (empty memory).	`None`
`delta_t`	`Tensor`	Must be `None` -- `DeltaMinGRU` never enables decay, so any non-`None` value raises `ValueError`.	`None`
`return_state`	`bool`	If True, also return the final state `h_T` (same flattened shape as `h_0`).	`False`

Returns:

Type	Description
`torch.Tensor or tuple of torch.Tensor`	`y`, shape `(B, T, hidden_size)`; or `(y, h_T)` when `return_state=True`.

Raises:

Type	Description
`ValueError`	If `delta_t` is not None, or `h_0`'s shape does not match `(B, 1, n_headsd_kd_v)`.

step ¶

step(x_t: Tensor, h_prev: Tensor | None = None, delta_t: Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]

Single recurrent step; same real-state convention as forward().

Computed from the same _coeffs helper forward() uses (applied once instead of over the full sequence). Unlike the other mixers, forward() does not also reuse this method's per-micro-step update (_micro_step) -- it runs the chunked-WY UT transform instead (_forward_chunked), a different but algebraically equivalent computation -- so forward/step parity here is guaranteed by the WY-transform derivation, not by shared code, and is pinned by the dual sequential/affine-scan oracle tests (TestDeltaMinGRU, TestMixerOracles).

Parameters:

Name	Type	Description	Default
`x_t`	`Tensor`	Input at the current timestep, shape `(B, input_size)`.	required
`h_prev`	`Tensor`	Previous per-head associative-memory state, flattened shape `(B, n_headsd_kd_v)`. Defaults to zero (empty memory).	`None`
`delta_t`	`Tensor`	Must be `None` -- `DeltaMinGRU` never enables decay, so any non-`None` value raises `ValueError`.	`None`

Returns:

Type	Description
`tuple of torch.Tensor`	`(y_t, h_t)`, shapes `(B, hidden_size)` and `(B, n_headsd_kd_v)` respectively. Unlike the other mixers' single-tensor `step` return, readout and state differ here (`y_t = H_t^T q_t`, `h_t` is the flattened `H_t`), so both are returned -- the honest generalization where readout != state.

Raises:

Type	Description
`ValueError`	If `delta_t` is not None, or `h_prev`'s shape does not match `(B, n_headsd_kd_v)`.

mingru.MinGRUBlock ¶

MinGRUBlock(d_model: int, mlp_expansion: int = 4, dropout: float = 0.0, learnable_h0: bool = False, mixer: str = 'log', mixer_kwargs: dict | None = None)

Bases: Module

Pre-norm residual block: LN -> minGRU -> +x, then LN -> MLP -> +x.

The minGRU is the (linear-in-state) sequence mixer; the MLP is the position-wise channel mixer. Because the scan's transition is diagonal, cross-channel interaction happens only here and in the next block's input projections.

Parameters:

Name	Type	Description	Default
`d_model`	`int`	Residual stream width (the minGRU maps `d_model -> d_model`).	required
`mlp_expansion`	`int`	Hidden width multiplier for the MLP. `0` disables the MLP sub-block entirely (scan-only block).	`4`
`dropout`	`float`	Applied after the minGRU output and inside the MLP.	`0.0`
`learnable_h0`	`bool`	Routed to the block's mixer when `mixer` is `"log"` or `"signed"` (see `MinGRU` / `SignedMinGRU`). Not accepted by `"rotation"`, `"givens"`, or `"delta"`: `RotationMinGRU`'s and `GivensMinGRU`'s `h_0` is an intrinsic learned parameter, and `DeltaMinGRU`'s is a fixed zero (empty associative memory) -- none of the three accept a `learnable_h0` flag, so this argument is silently unused for them.	`False`
`mixer`	`(log, signed, rotation, givens, delta)`	Selects the sequence mixer: `MinGRU` (log-space parallel scan), `SignedMinGRU` (signed diagonal transitions), `RotationMinGRU` (2x2 block rotations), `GivensMinGRU` (k-dim block rotations built from brick-wall Givens rounds), or `DeltaMinGRU` (DeltaNet/DeltaProduct generalized-Householder delta rule -- a per-head matrix state, not a diagonal/rotation scan; see its class docstring and the "Notes" section below). Any other value raises `ValueError`.	`"log"`
`mixer_kwargs`	`dict`	Extra constructor kwargs forwarded to the selected mixer class (e.g. `{"coupled": True}` for `"signed"`, `{"snap": (2, 3, 5)}` for `"rotation"`, `{"block_size": 8, "rounds": 3}` for `"givens"`, or `{"n_heads": 4, "nh": 2}` for `"delta"`); pass `decay`, `decay_rate`, `log1p_delta` here to enable time decay (see `MinGRU`/`SignedMinGRU`/`RotationMinGRU`/`GivensMinGRU`; `DeltaMinGRU` rejects non-`None` `decay`).	`None`

Notes

delta_t threads through forward/step straight to the mixer, unchanged; the block performs no decay-mode validation of its own. The mixer's own decay/delta_t pairing contract (ValueError both directions, see _validate_delta_t_pairing) is what fires whether the block is used standalone or inside a MinGRUStack.

Matrix-state mixers (carries_matrix_state = True -- currently only "delta") are called as forward(..., return_state=True) and their returned state is passed through as this block's state, instead of the output[:, -1:] convention every other mixer uses; this is the only behavior difference such a mixer causes, and it is why the returned state's last dim can differ from d_model for "delta" (it is n_heads*d_k*d_v, not d_model).

Examples:

>>> import torch
>>> from mingru import MinGRUBlock
>>> block = MinGRUBlock(d_model=64, mixer="signed")
>>> x = torch.randn(4, 128, 64)
>>> out, state = block(x)              # always returns (output, state)
>>> tuple(out.shape), tuple(state.shape)
((4, 128, 64), (4, 1, 64))

forward ¶

forward(x: Tensor, h_0: Tensor | None = None, delta_t: Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]

Parallel forward over a full sequence.

Always returns (output, state), matching step() and the nn.GRU convention; discard the state when not carrying it.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Residual-stream input, shape `(B, T, d_model)`.	required
`h_0`	`Tensor`	This block's mixer state carried from a previous chunk, shape `(B, 1, d_model)` for every mixer except `"delta"` (see `MinGRU.forward`); for `"delta"`, shape `(B, 1, n_headsd_kd_v)` (see `DeltaMinGRU.forward`).	`None`
`delta_t`	`Tensor`	Time gaps preceding each event, shape `(B, T)` or `(B, T, 1)`; forwarded to this block's mixer unchanged. Required iff the mixer's `decay` is enabled — the mixer's own pairing contract raises `ValueError` otherwise (see `_validate_delta_t_pairing`); this method adds no validation of its own.	`None`

Returns:

Type	Description
`tuple of torch.Tensor`	The output `(B, T, d_model)` and the block's final mixer state -- `(B, 1, d_model)` for every mixer except `"delta"`, `(B, 1, n_headsd_kd_v)` for `"delta"` (see the class "Notes" section) -- for the next chunk.

Raises:

Type	Description
`ValueError`	If the mixer's `decay` is enabled without `delta_t`, or `delta_t` is given without decay enabled.

step ¶

step(x_t: Tensor, h_prev: Tensor | None, delta_t: Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]

Single streaming step.

Parameters:

Name	Type	Description	Default
`x_t`	`Tensor`	Residual-stream input at time `t`, shape `(B, d_model)`.	required
`h_prev`	`Tensor or None`	This block's mixer state from `t-1`, shape `(B, d_model)` for every mixer except `"delta"` (see `MinGRU.step`); for `"delta"`, shape `(B, n_headsd_kd_v)` (see `DeltaMinGRU.step`). None at the first step.	required
`delta_t`	`Tensor`	Time gap preceding this event, shape `(B,)` or `(B, 1)`; forwarded to this block's mixer unchanged. Same required-iff-decay-enabled contract as `forward`.	`None`

Returns:

Type	Description
`tuple of torch.Tensor`	The block output, shape `(B, d_model)`, and the new mixer state -- `(B, d_model)` for every mixer except `"delta"`, `(B, n_headsd_kd_v)` for `"delta"`.

Raises:

Type	Description
`ValueError`	If the mixer's `decay` is enabled without `delta_t`, or `delta_t` is given without decay enabled.

mingru.MinGRUStack ¶

MinGRUStack(input_size: int, d_model: int, n_layers: int, mlp_expansion: int = 4, dropout: float = 0.0, learnable_h0: bool = False, mixer: str | list[str] = 'log', mixer_kwargs: dict | None = None, decay_layers: str = 'all')

Bases: Module

N stacked MinGRUBlocks with input projection and final norm.

Parameters:

Name	Type	Description	Default
`input_size`	`int`	Dimensionality of the raw inputs `x_t`.	required
`d_model`	`int`	Residual stream width.	required
`n_layers`	`int`	Number of blocks.	required
`mlp_expansion`	`int`	Per-block MLP expansion; `0` gives scan-only blocks.	`4`
`dropout`	`float`	Per-block dropout.	`0.0`
`learnable_h0`	`bool`	Passed through to every block; routed to the block's mixer when `mixer` is `"log"` or `"signed"`, unused for `"rotation"`, `"givens"`, and `"delta"` (see `MinGRUBlock`).	`False`
`mixer`	`str or list of str`	Sequence mixer for the blocks; one of `{"log", "signed", "rotation", "givens", "delta"}` (see `MinGRUBlock`). A single `str` uses that mixer for every block, bit-identical to prior versions: same construction order, same RNG consumption, same state_dict keys. Unknown `str` values raise `ValueError` (from `MinGRUBlock`, unchanged). A `list[str]` of length `n_layers` gives one mixer name per block, in order, so a single `"rotation"` block can live inside a deeper stack of `"signed"`/`"log"` blocks; a length mismatch or an unknown name anywhere in the list raises `ValueError` naming the valid set. More than one `"rotation"` entry emits exactly one `UserWarning` per construction (see Notes) and then proceeds; a single `"rotation"` entry in a mixed stack does not warn.	`"log"`
`mixer_kwargs`	`dict`	Schema is decided by the type of `mixer`, never by inspecting `mixer_kwargs` itself: `mixer: str` -> the current flat dict, forwarded as-is to every block's mixer constructor (e.g. `{"coupled": True}` for `"signed"`). `mixer: list[str]` -> `None`, or a dict keyed by mixer type name, e.g. `{"signed": {"coupled": True}, "rotation": {"snap": (2, 3, 6)}}`; each value is that type's flat kwargs dict, applied to every block of that type (two blocks of the same type share one config -- no per-block-index overrides). A key that is not a valid mixer name, or names a type absent from `mixer`, raises `ValueError`. A flat dict alongside a list `mixer`, or a type-keyed dict alongside a str `mixer`, raises `ValueError` describing both schemas.	`None`
`decay_layers`	`(all, last)`	Which blocks receive the decay-related keys (`decay`, `decay_rate`, `log1p_delta`) from each block's RESOLVED `mixer_kwargs` (the type-specific dict when `mixer` is a list). `"all"` (default, uniform decay) applies them to every block unchanged. `"last"` strips those three keys, by position, from every block except the final one (position `n_layers - 1`) -- whatever that block's mixer type. Any other value raises `ValueError` at construction. If a block's resolved kwargs carry no decay keys, `"last"` is a harmless no-op for it (nothing to strip). Trap in mixed stacks: `decay_layers` is purely positional, so under `"last"` the final block keeps its decay kwargs whatever type it is -- if that happens to be the `"rotation"` block, decay lands there and every `"signed"`/`"log"` block is stripped, regardless of where you put the decay keys in `mixer_kwargs`. Prefer per-type `mixer_kwargs` (put decay keys only under the type(s) you want decayed) to place decay deliberately in a mixed stack; reserve `decay_layers="last"` for homogeneous (single-`str`-mixer) stacks.	`"all"`

Notes

Shapes: forward maps (B, T, input_size) to (B, T, d_model); step maps (B, input_size) and a state (a list of n_layers tensors, one per block) to the output (B, d_model) and the updated state. Per-block state entries are opaque to MinGRUStack -- init_state() returns None placeholders (no shape assumed) and step() only ever passes a state entry through to its own block's .step(), never inspects or reshapes it -- so this is uniform across mixer types with one caveat: the entry's shape is (B, d_model) for every mixer except "delta" (carries_matrix_state = True), whose entry is (B, n_heads*d_k*d_v) (see MinGRUBlock.step). Both forward/step also accept an optional delta_t (same shapes as the mixers' delta_t): it is routed to a given block only if that block's mixer has decay enabled (checked via the mixer's own decay attribute), so a decay_layers="last" stack silently gives delta_t only to its final block. If NO block in the stack has decay enabled, passing delta_t raises ValueError (the mode-error rule); if at least one block decays, that block's own pairing contract raises if delta_t was left out.

More than one "rotation" block in mixer triggers exactly one UserWarning per construction: the straight-through snap discontinuity can compound across rotation layers, and multi- rotation stacks are validated only at L=2 on the S3 probe (deeper is untested — see the README's Rotation variant section, "Depth: L=2 is validated"); construction proceeds regardless. A single "rotation" block in a mixed stack does not warn. Multiple "givens" blocks do NOT warn: the warning is specific to the straight-through snap discontinuity, and GivensMinGRU is continuous (unsnapped), so it has no such discontinuity to compound.

Examples:

>>> import torch
>>> from mingru import MinGRUStack
>>> model = MinGRUStack(input_size=32, d_model=64, n_layers=3, mixer="log")
>>> x = torch.randn(4, 128, 32)
>>> out, state = model(x)             # state: one entry per block
>>> tuple(out.shape), len(state)
((4, 128, 64), 3)

forward ¶

forward(x: Tensor, state: list[Tensor | None] | None = None, delta_t: Tensor | None = None) -> tuple[torch.Tensor, list[torch.Tensor]]

Parallel training-mode forward.

Always returns (output, state), matching step() and the nn.GRU convention; discard the state when not carrying it.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input sequence, shape `(B, T, input_size)`.	required
`state`	`list of (torch.Tensor or None)`	Per-block minGRU states from a previous chunk — as returned by a prior `forward()`, or from streaming `step()` after unsqueezing each entry to `(B, 1, d_model)`. For TBPTT, detach the returned state before feeding it to the next chunk.	`None`
`delta_t`	`Tensor`	Time gaps preceding each event, shape `(B, T)` or `(B, T, 1)`. Routed only to blocks whose mixer has decay enabled (see class Notes); required iff at least one block decays, and rejected with `ValueError` if no block decays.	`None`

Returns:

Type	Description
`tuple`	The output `(B, T, d_model)` and a list of `n_layers` per-block final states, each `(B, 1, d_model)` — except matrix-state blocks (`mixer="delta"`), whose entry is the flattened `(B, 1, n_headsd_kd_v)` state; treat entries as opaque.

Raises:

Type	Description
`ValueError`	If `delta_t` is given but no block's mixer has decay enabled, or a decayed block's mixer requires `delta_t` and it was omitted.

init_state ¶

init_state() -> list[None]

Fresh streaming state.

Returns:

Type	Description
`list of None`	One slot per block; pass to `step()` at the first timestep.

step ¶

step(x_t: Tensor, state: list[Tensor | None], delta_t: Tensor | None = None) -> tuple[torch.Tensor, list[torch.Tensor]]

Streaming step.

Total cached state is n_layers * d_model per sample, except that a "delta" block's entry is n_heads*d_k*d_v instead of d_model (see the class "Notes" section).

Parameters:

Name	Type	Description	Default
`x_t`	`Tensor`	Input at the current timestep, shape `(B, input_size)`.	required
`state`	`list of (torch.Tensor or None)`	From `init_state()` or a previous `step()`. Each entry is opaque to this method -- passed straight through to its own block's `.step()`, never inspected or reshaped here.	required
`delta_t`	`Tensor`	Time gap preceding this event, shape `(B,)` or `(B, 1)`. Same per-block routing and `ValueError` contract as `forward`.	`None`

Returns:

Type	Description
`tuple`	The output, shape `(B, d_model)`, and the updated state (a list of `n_layers` tensors, one per block -- see the class "Notes" section for each entry's shape).

Raises:

Type	Description
`ValueError`	If `delta_t` is given but no block's mixer has decay enabled, or a decayed block's mixer requires `delta_t` and it was omitted.