Scan operations¶

The four associative scan primitives the mixers are built on. Each accepts MINGRU_SCAN dispatch (see Control scan dispatch); the eager signatures and shape contracts are generated from the source docstrings.

mingru.parallel_scan_log ¶

parallel_scan_log(log_coeffs: Tensor, log_values: Tensor) -> torch.Tensor

Heinsen (2023) parallel scan in log-space.

Solves h_t = a_t * h_{t-1} + b_t for all t simultaneously, given log(a_t) and log(b_t), assuming a_t, b_t > 0.

Parameters:

Name	Type	Description	Default
`log_coeffs`	`Tensor`	Shape `(B, T, D)`. `log(a_t)` for `t = 1..T`.	required
`log_values`	`Tensor`	Shape `(B, T + 1, D)`. Slot 0 is `log(h_0)` — the scan's initial value, i.e. `v_0 <- h_0` in Heinsen's formulation; slots `1..T` are `log(b_t)`.	required

Returns:

Type	Description
`Tensor`	Shape `(B, T, D)`. The states `h_1..h_T`.

References

Heinsen (2023), "Efficient Parallelization of a Ubiquitous Sequential Computation" -- the log-space cumulative-sum / logcumsumexp scan.

Notes

Under MINGRU_SCAN=auto (the default) on CUDA, this op dispatches to a Triton kernel that is measured SLOWER than eager for a full forward+ backward training step at benchmarked shapes (0.70-0.80x eager; see experiments/bench/scan_bench.md), even though its forward-only latency is roughly on par with eager. MINGRU_SCAN=eager is the right choice for training MinGRU/SignedMinGRU today.

Examples:

>>> import torch
>>> from mingru import parallel_scan_log
>>> B, T, D = 2, 5, 3
>>> log_coeffs = -torch.rand(B, T, D)          # log(a_t), a_t in (0, 1)
>>> log_values = -torch.rand(B, T + 1, D)      # log(h_0) and log(b_t)
>>> h = parallel_scan_log(log_coeffs, log_values)
>>> tuple(h.shape)
(2, 5, 3)

mingru.linear_scan ¶

linear_scan(a: Tensor, b: Tensor) -> tuple[torch.Tensor, torch.Tensor]

Linear-space associative scan for h_t = a_t * h_{t-1} + b_t.

Hillis-Steele doubling over the segment-composition monoid (A1, B1) o (A2, B2) = (A2*A1, A2*B1 + B2): O(T log T) work, O(log T) depth, pure torch ops, differentiable. Unlike parallel_scan_log, coefficients may be NEGATIVE — this is what enables signed transition eigenvalues (see SignedMinGRU).

Parameters:

Name	Type	Description	Default
`a`	`Tensor`	Shape `(B, T, D)`. Transition coefficients `a_t`; stability expects `\|a_t\| <= 1` but is not enforced.	required
`b`	`Tensor`	Shape `(B, T, D)`. Additive inputs `b_t`.	required

Returns:

Type	Description
`tuple of torch.Tensor`	`(A, Bc)`, each `(B, T, D)`, where `h_t = A_t * h_0 + Bc_t` with `A_t` the running product of `a` and `Bc_t` the `h_0 = 0` solution of the recurrence.

Notes

Hillis-Steele is work-inefficient (O(T log T) vs O(T) for a work-efficient Blelloch scan) and retains O(log T) full (B, T, D) tensors for autograd. That is a deliberate simplicity-over-efficiency choice: the overhead is negligible at the short training lengths this repo targets (probes train at T=64), and no stable torch.associative_scan primitive exists to lean on. Revisit for long-sequence training.

Examples:

>>> import torch
>>> from mingru import linear_scan
>>> a = torch.rand(2, 5, 3) * 2 - 1        # a_t in (-1, 1), signed
>>> b = torch.randn(2, 5, 3)
>>> A, Bc = linear_scan(a, b)
>>> tuple(A.shape), tuple(Bc.shape)
((2, 5, 3), (2, 5, 3))

mingru.matrix_scan ¶

matrix_scan(M: Tensor, b: Tensor) -> tuple[torch.Tensor, torch.Tensor]

Associative scan for h_t = M_t @ h_{t-1} + b_t, 2x2 block transitions.

The non-commutative generalization of linear_scan: transitions are 2x2 matrices composed by matrix multiplication instead of scalars multiplied together, so the running product depends on order. This is what lets a mixer built on matrix_scan (RotationMinGRU) express non-abelian state tracking that a diagonal/commutative scan (linear_scan, parallel_scan_log) provably cannot. Same Hillis-Steele doubling scheme as linear_scan, over the affine 2x2 monoid (A1, B1) o (A2, B2) = (A2 @ A1, A2 @ B1 + B2): O(T log T) work, O(log T) depth, pure torch ops, differentiable. The identity padding is the 2x2 identity matrix (not 1.0), and composition order is A_current @ A_earlier — matrix multiplication does not commute, unlike linear_scan's scalar product.

Parameters:

Name	Type	Description	Default
`M`	`Tensor`	Shape `(B, T, n, 2, 2)`. Per-block transition matrices `M_t`; stability expects `M_t`'s spectral norm `<= 1` but this is not enforced.	required
`b`	`Tensor`	Shape `(B, T, n, 2)`. Additive inputs `b_t`, viewed as `n` 2-vectors per timestep.	required

Returns:

Type	Description
`tuple of torch.Tensor`	`(A, Bc)`, shapes `(B, T, n, 2, 2)` and `(B, T, n, 2)`, each aligned to `t`, where `h_t = A_t @ h_0 + Bc_t` with `A_t` the running matrix product of `M` and `Bc_t` the `h_0 = 0` solution of the recurrence.

Notes

Same simplicity-over-efficiency tradeoff as linear_scan: Hillis-Steele is work-inefficient (O(T log T) vs O(T)) and retains O(log T) full (B, T, n, 2, 2) tensors for autograd. Fine at the short training lengths this repo targets (probes train at T=64); revisit for long-sequence training.

Examples:

>>> import torch
>>> from mingru import matrix_scan
>>> B, T, n = 2, 5, 4
>>> M = torch.eye(2).expand(B, T, n, 2, 2)     # identity transitions
>>> b = torch.randn(B, T, n, 2)
>>> A, Bc = matrix_scan(M, b)
>>> tuple(A.shape), tuple(Bc.shape)
((2, 5, 4, 2, 2), (2, 5, 4, 2))

mingru.matrix_affine_scan ¶

matrix_affine_scan(A: Tensor, Bm: Tensor) -> tuple[torch.Tensor, torch.Tensor]

Associative scan for H_t = A_t @ H_{t-1} + B_t, k x k transitions.

The k x k generalization of matrix_scan: transitions are square k x k matrices (instead of matrix_scan's fixed 2x2 blocks) and injections are rectangular k x v matrix-valued states (instead of matrix_scan's 2-vectors, i.e. v = 1). Both scan the same affine monoid (A1, B1) o (A2, B2) = (A2 @ A1, A2 @ B1 + B2) by the same Hillis-Steele doubling scheme; matrix_scan is exactly the k = 2, v = 1 special case, and matrix_affine_scan is what the k > 2 block-rotation mixer (GivensMinGRU, with the state as a v = 1 column) and matrix-valued delta-rule states run on. O(T log T) work, O(log T) depth, pure torch ops, differentiable. Composition order is A_current @ A_earlier — matrix multiplication does not commute — so, as in matrix_scan, the running product depends on order.

Parameters:

Name	Type	Description	Default
`A`	`Tensor`	Shape `(B, T, n, k, k)`. Per-block square transitions `A_t`; stability expects each `A_t`'s spectral norm `<= 1` but this is not enforced.	required
`Bm`	`Tensor`	Shape `(B, T, n, k, v)`. Additive injections `B_t`, viewed as `n` `k x v` matrices per timestep (`v = 1` recovers the vector state of `matrix_scan`).	required

Returns:

Type	Description
`tuple of torch.Tensor`	`(Abar, Bbar)`, shapes `(B, T, n, k, k)` and `(B, T, n, k, v)`, each aligned to `t`, where `H_t = Abar_t @ H_0 + Bbar_t` with `Abar_t` the running matrix product of `A` and `Bbar_t` the `H_0 = 0` solution of the recurrence.

Notes

Same simplicity-over-efficiency tradeoff as matrix_scan / linear_scan: Hillis-Steele is work-inefficient (O(T log T) vs O(T)) and retains O(log T) full (B, T, n, k, k) tensors for autograd. Fine at the short training lengths this repo targets (probes train at T=64); revisit for long-sequence training. Kept as a distinct helper beside matrix_scan (rather than a generalization of it) so the recorded 2x2 rotation evidence's floating-point path is preserved byte-for-byte.

Examples:

>>> import torch
>>> from mingru import matrix_affine_scan
>>> B, T, n, k, v = 2, 5, 4, 8, 1
>>> A = torch.eye(k).expand(B, T, n, k, k)     # identity transitions
>>> Bm = torch.randn(B, T, n, k, v)
>>> Abar, Bbar = matrix_affine_scan(A, Bm)
>>> tuple(Abar.shape), tuple(Bbar.shape)
((2, 5, 4, 8, 8), (2, 5, 4, 8, 1))