1. Sequential Linear-Recurrence Scan

Problem

Real-time voice models built on state-space models (SSMs) — the architecture family Cartesia builds on — process audio one frame at a time through a per-channel linear recurrence: h_t = a_t * h_{t-1} + b_t, where a_t is a (typically input-dependent) decay gate and b_t is the input-driven increment at that timestep. This single elementwise recurrence is the computational core shared by S4, S4D, Mamba's selective scan, and GateLoop's data-controlled linear recurrence — everything else in those architectures (convolutions, gating projections, output heads) wraps around this primitive.

Before reasoning about how to make this recurrence fast (parallel scan, chunking — the next two problems), you need the reference implementation cold: a plain sequential loop over the time dimension that is unambiguously correct, even though it can't exploit a GPU's parallelism because each timestep depends on the previous one's output.

Source: src/1_sequential_linear_recurrence_scan.py

def sequential_scan(a: torch.Tensor, b: torch.Tensor, h0: torch.Tensor | None = None) -> torch.Tensor:
    ...

Examples:
>>> a = torch.tensor([[[0.5], [0.5], [0.5]]])
>>> b = torch.tensor([[[1.0], [1.0], [1.0]]])
>>> sequential_scan(a, b)
tensor([[[1.0000], [1.5000], [1.7500]]])

Step-by-Step Approach

1. Read off the shapes: a and b are (batch, length, dim) — the recurrence runs independently per batch element and per channel (dim), with no mixing across channels.
2. Initialize the running hidden state h_prev to h0 if given, else a zero tensor of shape (batch, dim).
3. Loop over the time dimension t = 0 .. length-1. At each step compute h_t = a[:, t, :] * h_prev + b[:, t, :], an elementwise (Hadamard) multiply-add, not a matrix multiply.
4. Append each h_t to a list and update h_prev = h_t for the next iteration.
5. After the loop, stack the per-timestep states along a new time axis (dim=1) to produce the full (batch, length, dim) output.
6. Sanity-check against a tiny hand-computable case: with a=0.5 and b=1.0 constant and h0=0, you should get h_1=1.0, h_2=0.5*1+1=1.5, h_3=0.5*1.5+1=1.75 — if your loop doesn't reproduce this by hand, the bug is almost always an off-by-one in which timestep's a/b you're using.

The key insight is that this recurrence is elementwise, not a matrix recurrence like a vanilla RNN's h_t = tanh(W h_{t-1} + ...) — each channel evolves independently with its own scalar gate, which is exactly what makes it cheap (O(dim) work per timestep instead of O(dim²)) and exactly what makes it parallelizable across channels and, as the next problem shows, across time too.

Reference solution

def sequential_scan(a: torch.Tensor, b: torch.Tensor, h0: torch.Tensor | None = None) -> torch.Tensor:
    batch, length, dim = a.shape
    if h0 is None:
        h0 = torch.zeros(batch, dim, dtype=a.dtype, device=a.device)
    h_prev = h0  # (batch, dim)
    outputs = []
    for t in range(length):
        # h_t = a_t * h_{t-1} + b_t, elementwise per (batch, dim)
        h_t = a[:, t, :] * h_prev + b[:, t, :]
        outputs.append(h_t)
        h_prev = h_t
    # stack along the time axis -> (batch, length, dim)
    return torch.stack(outputs, dim=1)

Key Functions & Tricks

• a[:, t, :] — slice out one timestep across all batches/channels, shape (batch, dim)
• torch.stack(outputs, dim=1) — stack a Python list of (batch, dim) tensors into (batch, length, dim)
• torch.zeros(batch, dim, dtype=..., device=...) — match dtype/device of the default initial state to the inputs
• Elementwise (Hadamard) recurrence — diagonal state transition, O(dim) per step, not a matrix-vector product
• h0: torch.Tensor | None = None — optional-initial-state pattern common in streaming/incremental APIs (carrying state across chunks or calls)

How to Recognize This Pattern