12. Chunk-wise Recurrent Scan with Carried State

Problem

A real-time voice server can't wait for an entire utterance before it starts producing audio — it processes fixed-size windows ("chunks") of frames as they arrive and must carry the recurrent state from the end of one chunk into the start of the next, so that chunk boundaries are invisible in the output: running the recurrence chunk-by-chunk must give bit-identical results to running it once over the whole sequence.

Given the recurrence h_t = a_t * h_{t-1} + x_t (elementwise), implement a chunked scan that processes chunk_size timesteps at a time, carries the final hidden state of each chunk into the next chunk's initial state, and handles a sequence length that is not an exact multiple of chunk_size (the last chunk is simply shorter).

Source: src/12_chunked_linear_scan.py

def chunked_linear_scan(
    x: torch.Tensor, a: torch.Tensor, h0: torch.Tensor, chunk_size: int,
) -> torch.Tensor: ...

>>> x = torch.randn(2, 7, 4)
>>> a = torch.sigmoid(torch.randn(2, 7, 4))
>>> h0 = torch.zeros(2, 4)
>>> chunked_linear_scan(x, a, h0, chunk_size=3).shape
torch.Size([2, 7, 4])

Step-by-Step Approach

1. Iterate over chunk start indices with a step of chunk_size: range(0, seq_len, chunk_size).
2. For each chunk, clamp its end index to min(start + chunk_size, seq_len) so the final, possibly-shorter chunk doesn't index out of bounds.
3. Within a chunk, run the ordinary sequential recurrence timestep by timestep, always reading the current carried hidden state h as the previous step's output.
4. After finishing a chunk, do not reset h — it must flow directly into the first timestep of the next chunk as that chunk's initial state, exactly as if no chunk boundary existed.
5. Collect every per-timestep hidden state (not just the per-chunk final ones) into the output, and stack them into a single (batch, seq_len, hidden) tensor at the end.
6. Validate against an unchunked reference: the same recurrence run as one single sequential pass over all seq_len steps must match exactly, for any chunk_size including one that doesn't evenly divide seq_len.

The key insight is that "chunking" only changes how the loop is organized (an outer loop over chunk boundaries, an inner loop over timesteps within a chunk) — it never changes the math, because the one piece of state that crosses a chunk boundary, h, is threaded through unchanged.

Reference solution

import torch


def chunked_linear_scan(
    x: torch.Tensor, a: torch.Tensor, h0: torch.Tensor, chunk_size: int,
) -> torch.Tensor:
    batch, seq_len, hidden = x.shape
    h = h0
    outputs = []
    # outer loop advances by whole chunks -- this is the boundary a
    # real-time server actually pays for (one chunk = one arrived window of
    # audio); the inner loop is the sequential recurrence within that chunk
    for start in range(0, seq_len, chunk_size):
        end = min(start + chunk_size, seq_len)
        for t in range(start, end):
            h = a[:, t, :] * h + x[:, t, :]
            outputs.append(h)
        # h now holds the state carried across the chunk boundary into the
        # next iteration's initial state -- nothing besides h survives
    return torch.stack(outputs, dim=1)

Key Functions & Tricks

• range(0, seq_len, chunk_size) — generates chunk start offsets without needing a separate counter or ceiling-division computation.
• min(start + chunk_size, seq_len) — clamps the last chunk's end index so a non-divisible sequence length never indexes past the tensor.
• Carrying a single scalar-per-batch-row Python variable h across both loop levels — the entire "carried state" API surface a streaming caller needs is this one variable's value at chunk boundaries.
• torch.stack(outputs, dim=1) — every timestep's hidden state is appended regardless of which chunk it came from, so the final shape is indistinguishable from an unchunked run.

How to Recognize This Pattern