3. Batching Variable-Length Sequences with Masked Recurrence

Problem

Real conversations don't arrive in neat, equal-length batches: one caller's utterance is 40 frames, another's is 400. Batch-scoring a recurrent/SSM layer efficiently means stacking these different-length sequences into a single padded tensor — but the padded positions in a short row contain garbage, and if the recurrence runs naively over them, that garbage propagates into the row's hidden state exactly like real input would. A correct implementation runs the recurrence for the whole batch at once while making padded timesteps a pure no-op for their row: the hidden state must simply freeze at each row's own last valid timestep.

There's a second trap beyond masking during the scan: extracting the "final" hidden state per row. Because rows have different lengths, h_all[:, -1, :] is only correct for the single longest row in the batch — every shorter row's real final state has to be gathered using that row's own length.

Source: src/3_masked_batched_recurrence.py

def masked_batched_recurrence(
    x: torch.Tensor, a: torch.Tensor, lengths: torch.Tensor, h0: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]: ...

>>> x = torch.randn(3, 4, 2)
>>> a = torch.sigmoid(torch.randn(3, 4, 2))
>>> lengths = torch.tensor([2, 4, 1])
>>> h0 = torch.zeros(3, 2)
>>> h_all, h_final = masked_batched_recurrence(x, a, lengths, h0)
>>> h_all.shape, h_final.shape
(torch.Size([3, 4, 2]), torch.Size([3, 2]))

Step-by-Step Approach

1. Loop over timesteps t from 0 to max_len - 1, running the recurrence for the whole batch at once at each step (not per-row) for GPU efficiency.
2. At each t, build a boolean mask valid = (t < lengths).unsqueeze(-1), shape (batch, 1), so it broadcasts across the hidden dimension.
3. Compute the candidate next state h_candidate = a[:, t, :] * h + x[:, t, :] as usual.
4. Use torch.where(valid, h_candidate, h) to select the candidate for still-valid rows and simply keep the previous h unchanged for rows already past their length — not an in-place masked assignment, which is easy to get backwards.
5. Stack every timestep's h into h_all, then gather each row's own final state using lengths - 1 as a per-row index rather than slicing the last column.
6. Use gather along the time dimension with an index tensor built from lengths - 1, expanded to match the hidden dimension, to pull out exactly one timestep per row in a single vectorized call.

The key insight is that masking a recurrence is not about zeroing outputs after the fact — it's about making the state-update itself a no-op past each row's true length, via torch.where choosing between "advance" and "hold," so garbage in the padded region can never leak into a real row's trajectory.

Reference solution

import torch


def masked_batched_recurrence(
    x: torch.Tensor, a: torch.Tensor, lengths: torch.Tensor, h0: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    batch, max_len, hidden = x.shape
    h = h0
    outputs = []
    for t in range(max_len):
        # valid[i] is True while row i still has real (non-padded) data at
        # this timestep; shape (batch, 1) so it broadcasts over hidden
        valid = (t < lengths).unsqueeze(-1)
        h_candidate = a[:, t, :] * h + x[:, t, :]
        # torch.where, not in-place masked assignment: rows past their own
        # length keep exactly their previous state (a true no-op), instead
        # of being advanced with whatever garbage sits in the padding
        h = torch.where(valid, h_candidate, h)
        outputs.append(h)
    h_all = torch.stack(outputs, dim=1)

    # gather each row's state at its own last valid index (lengths - 1),
    # not simply h_all[:, -1, :] which is only right for the longest row
    gather_idx = (lengths - 1).view(batch, 1, 1).expand(batch, 1, hidden)
    h_final = h_all.gather(dim=1, index=gather_idx).squeeze(1)
    return h_all, h_final

Key Functions & Tricks

• torch.where(valid, h_candidate, h) — the core masking trick: selects between the advanced state and the held-over state per row, fully vectorized, no Python-level branching per row.
• (t < lengths).unsqueeze(-1) — broadcasts a per-row boolean comparison against a scalar timestep into the shape needed to gate a (batch, hidden) tensor.
• tensor.gather(dim=1, index=...) — pulls one time-index per row out of the padded (batch, max_len, hidden) tensor in a single call, avoiding a Python loop over rows.
• (lengths - 1).view(batch, 1, 1).expand(batch, 1, hidden) — reshapes and broadcasts the per-row target index to the exact shape gather requires.

How to Recognize This Pattern