nn.scan over an RNN cell

Apply one GRU cell across a sequence

Use Linen's nn.scan to turn a single-step nn.GRUCell into a sequence operation. Start from a zero hidden state, read one input row per step and return the final hidden state.

The input has shape (T, D). A cell call takes a carry of shape (H,) and one row of shape (D,), then returns a new carry and an output, each with shape (H,). For a GRU, the new carry and output are equal.

zero state -> cell(x[0]) -> cell(x[1]) -> ... -> cell(x[T-1])
                 |              |                    |
                ys[0]          ys[1]                ys[T-1]
                                                     |
                                                final carry

All steps reuse one set of GRU parameters. There is no batch dimension in this exercise and no extra transformation before the first input row.

Tell the lifted scan what to share

Construct the transformed module class with:

ScanGRU = nn.scan(
    nn.GRUCell,
    variable_broadcast="params",
    split_rngs={"params": False},
    in_axes=0,
    out_axes=0,
)

These options do different jobs:

• variable_broadcast="params" makes the parameter collection shared across iterations. It contains the GRU's input and recurrent matrices and biases, not just one weight matrix.
• split_rngs={"params": False} keeps the parameter-initialization random stream unsplit along the scan axis. It does not decide whether parameter arrays have a time axis; the variable options do that.
• in_axes=0 reads rows from x. out_axes=0 stacks the cell outputs into ys with shape (T, H). Neither option adds a time axis to the carry, which remains (H,).

Per-step parameter arrays require a different setup, such as variable_axes={"params": 0}. Simply omitting variable_broadcast does not automatically create independent weights. Likewise, splitting an RNG is not an instruction to create or flatten parameter copies.

Carry and output history

Calling the scanned cell returns (final_carry, ys). For a nonempty sequence, the GRU's final carry equals ys[-1]. This exercise returns only that carry, flattened to length H.

Here is a complete example that checks the scanned call against a sequence of ordinary cell calls using exactly the same initialized parameters:

import jax
import jax.numpy as jnp
import flax.linen as nn

ScanGRU = nn.scan(
    nn.GRUCell, variable_broadcast="params",
    split_rngs={"params": False}, in_axes=0, out_axes=0,
)
x = jnp.array([[.1, .2], [.3, -.1], [0., .4]])
initial = jnp.zeros(2)
scanned = ScanGRU(features=2)
variables = scanned.init(jax.random.PRNGKey(0), initial, x)
final, ys = scanned.apply(variables, initial, x)

cell = nn.GRUCell(features=2)
carry = initial
outputs = []
for row in x:
    carry, y = cell.apply(variables, carry, row)
    outputs.append(y)

assert ys.shape == (3, 2)
assert jnp.allclose(ys, jnp.stack(outputs), atol=1e-6)
assert jnp.allclose(final, carry, atol=1e-6)
assert jnp.allclose(final, ys[-1], atol=1e-6)

This comparison checks parameter sharing and time order. It is not a speed benchmark. Inside JIT, a Python loop with a fixed number of iterations is unrolled into the traced computation; scan represents the recurrence as a loop. That can keep the compiled program smaller, but runtime speed depends on the workload.

Why use a lifted transformation?

nn.scan handles Linen variable collections and random streams according to the options you provide. It does not automatically choose the right sharing policy for every module. Carry structure, shape and dtype must stay fixed across steps.

A pure function calling cell.apply with an already initialized, unchanged variable dictionary can also be used inside jax.lax.scan; those variables can be closed over rather than threaded as carry. The lifted form here lets the scanned module initialize and manage its parameters through the Linen interface.

Pass both the initial carry and the sequence to init. Avoid reasoning about initialization as exactly one ordinary timestep: parameter setup and tracing are implementation details, while apply defines the recurrence whose outputs this exercise checks.