15. Selective (Data-Dependent) SSM Scan

Problem

Classic linear state-space layers (S4) use a fixed, input-independent transition: the same decay applies at every timestep regardless of what the input actually is. Mamba's key change is making the recurrence selective: the decay, input-projection, and output-projection terms are all themselves functions of the current input, so the model can choose — per timestep, per channel — to let information flow through or to reset it. That selectivity is what gives these models content-aware behavior close to attention while keeping a constant-size recurrent state and linear-time inference.

This problem asks for the core recurrence itself, with the data-dependent parameters already computed and handed to you (in the real architecture they come from a small input-conditioned linear projection; here they're just given tensors so the problem is squarely about the scan, not the projection).

Source: src/15_selective_ssm_scan.py

def selective_ssm_scan(
    x: torch.Tensor, A: torch.Tensor, B: torch.Tensor,
    C: torch.Tensor, h0: torch.Tensor,
) -> torch.Tensor: ...

>>> batch, seq_len, d = 2, 6, 4
>>> x, B, C = torch.randn(batch, seq_len, d), torch.randn(batch, seq_len, d), torch.randn(batch, seq_len, d)
>>> A = torch.sigmoid(torch.randn(batch, seq_len, d))
>>> h0 = torch.zeros(batch, d)
>>> selective_ssm_scan(x, A, B, C, h0).shape
torch.Size([2, 6, 4])

Step-by-Step Approach

1. Recognize the recurrence's two stages per timestep: a state update h_t = A_t * h_{t-1} + B_t * x_t, then a readout y_t = C_t * h_t — the state and the output are two separate quantities.
2. Loop over timesteps, indexing A[:, t, :], B[:, t, :], C[:, t, :], and x[:, t, :] at each step — every one of these is read fresh per timestep, which is exactly the "selective" part (a classic S4 layer would instead use one fixed A shared across all timesteps).
3. Update the state elementwise: h = A_t * h + B_t * x_t. All of A, B, C, x, h share the same trailing dimension d, so every operation here is elementwise, not a matmul.
4. Compute the readout y_t = C_t * h from the updated state, not the previous one.
5. Collect every timestep's y_t and stack them into the final (batch, seq_len, d) output sequence.
6. Sanity-check the degenerate case A = 0 everywhere: the state should then depend only on the current timestep's B and x, never on history — a useful way to confirm you haven't accidentally left a stale h_prev term in the update.

The key insight is that "selective" doesn't change the shape of the recurrence at all versus a plain gated linear recurrence — it only changes where A, B, and C come from (a per-token projection instead of fixed learned parameters), so the scan itself is exactly the same elementwise update-then-readout loop.

Reference solution

import torch


def selective_ssm_scan(
    x: torch.Tensor, A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, h0: torch.Tensor,
) -> torch.Tensor:
    batch, seq_len, d = x.shape
    h = h0
    ys = []
    for t in range(seq_len):
        # every one of A, B, C is read at this specific timestep -- this is
        # the "selective" part: the transition itself changes per token,
        # unlike a classic S4 layer's fixed, input-independent A
        h = A[:, t, :] * h + B[:, t, :] * x[:, t, :]
        y_t = C[:, t, :] * h
        ys.append(y_t)
    return torch.stack(ys, dim=1)

Key Functions & Tricks

• Elementwise * throughout, never @/matmul — because A, B, C are per-channel diagonal transitions here, not dense matrices, the whole recurrence is a Hadamard product plus an add at every step.
• A[:, t, :]-style slicing — pulls out the timestep-specific parameters, shape (batch, d), keeping the loop body free of any dependence on seq_len beyond the index.
• torch.stack(ys, dim=1) — assembles the per-timestep readouts into the (batch, seq_len, d) output, inserting the sequence axis at position 1.
• Reading y_t from the state after the update, not before — a one-line ordering detail that's easy to get backwards under interview pressure.

How to Recognize This Pattern