vmap over lax.scan

Run the same recurrence on independent rows

Each row of a batch may represent a sequence: a sentence, a signal, or a series of observations. A scan handles the dependency between successive values within one row. A vmap applies that row-level calculation across the batch.

Write a helper that scans one row, then pass the helper to jax.vmap. Each row has its own logical carry. Updating one row never adds to another row's total. Vmap transforms the operations to handle a batch; it does not launch a separate Python loop for every row or guarantee a particular hardware speedup.

Worked mini-example

import jax
import jax.numpy as jnp
from jax import lax

def step(total, value):
    total += value
    return total, total

def per_row(row):
    _, sums = lax.scan(
        step, 0.0, row
    )
    return sums

x_batch = jnp.array([
    [1.0, 2.0, 3.0],
    [4.0, 5.0, 6.0],
])
batched = jax.vmap(per_row, in_axes=0, out_axes=0)
result = batched(x_batch)
# [[1.0, 3.0, 6.0],
#  [4.0, 9.0, 15.0]]

lax.scan(step, init, xs) returns (final_carry, stacked_outputs). The second element contains the stacked y values emitted at each step.

Common pitfalls

• Passing a result instead of a function. vmap needs a callable; lax.scan(...) returns a tuple. The per_row helper keeps the scan body and initial carry fixed while exposing the row to be mapped. A suitably configured partial function is another option.
• Changing the carry's shape or dtype. Every step must return a carry matching the initial carry's shape and dtype. This example uses floating point inputs and a scalar carry. jnp.zeros((), dtype=row.dtype) makes that carry's dtype match the row explicitly.
• Confusing batch and time. The outer map selects rows; scan visits values within each row in order. The next row starts again from zero.
• Forgetting to unpack outputs. lax.scan returns (final_carry, ys). Discard the carry with _ when you only need the per-step outputs.

Follow the two axes

For an input with shape (N, T), axis 0 selects one of N rows. The helper receives shape (T,); its scan visits those T elements. Each helper returns T prefix sums. With out_axes=0, vmap places the batch dimension first again, producing shape (N, T).

In the example, the first row's carry progresses 0 → 1 → 3 → 6. The second starts afresh: 0 → 4 → 9 → 15. Flattening the input before scanning would incorrectly carry the first row's total into the second. Scanning the whole matrix directly instead visits rows, which is also a different axis choice.

This diagram shows the outer map's shapes, not the sequential steps inside a row: