NNX Scan RNN

Read a sequence with one recurrent cell

A recurrent network updates a hidden vector as it reads input rows in order. At step t, a GRU cell combines the previous hidden state with x[t] and returns the next hidden state. The same cell parameters are used at every step; they are not trained or reinitialized during this forward pass.

jax.lax.scan expresses that recurrence with a body of the form (carry, x_t) -> (new_carry, y_t). Here carry has shape (H,) and x_t has shape (D,). NNX's GRUCell returns the new hidden state as both new_carry and y_t.

After T steps, scan returns a final carry of shape (H,) and stacked outputs of shape (T, H). This task returns only the final carry. It represents the state after reading the requested prefix, not a probability distribution or a prediction from a trained model.

A complete two-step example

The following example sets simple parameters so you can follow the hidden state. Zero gate inputs give an update gate of 0.5. A candidate bias of arctanh(0.5) gives a candidate value of 0.5. Each step therefore computes new_h = 0.5 * 0.5 + 0.5 * h: zero becomes 0.25, then 0.375.

import jax
import jax.numpy as jnp
from flax import nnx

cell = nnx.GRUCell(1, 1, rngs=nnx.Rngs(0))
cell.dense_i.kernel[...] = jnp.zeros((1, 3))
cell.dense_h.kernel[...] = jnp.zeros((1, 3))
cell.dense_i.bias[...] = jnp.array([0., 0., jnp.arctanh(0.5)])

def step(carry, x_t):
    return cell(carry, x_t)

x = jnp.zeros((2, 1))
final, outputs = jax.lax.scan(step, jnp.zeros((1,)), x)
assert jnp.allclose(outputs, jnp.array([[0.25], [0.375]]))
assert jnp.allclose(final, jnp.array([0.375]))

For the exercise, retain the GRUCell's default initialized parameters; do not replace them with the example's known values.

What happens during tracing

A Python loop inside a JIT-compiled function can produce a computation graph with a copy of the cell operations for every step. Scan represents the body as a loop, which can reduce compilation work for long sequences. It does not promise a particular speedup or automatically unroll the entire loop. Runtime still includes the requested sequence of updates.

In this recipe, step closes over a cell built outside the loop. During tracing, its operations read the cell's parameter arrays. The compiled computation does not look up Python attributes again at every execution. In particular, changing an external captured cell after compiling a function is not a reliable way to update that function's parameters.

This read-only GRU forward works with an array carry and plain lax.scan. It is not a general recipe for mutating an NNX module or consuming a mutable RNG stream inside a raw JAX transformation. Use NNX-aware transformations and explicit state handling for those cases.

If you later differentiate a loss with respect to cell parameters, pass the cell explicitly to nnx.value_and_grad and define the scan inside that loss. The inner closure then reads the parameters supplied to the differentiated function. Closing over a cell does not, by itself, request gradients for it.

Keep the carry and time axis consistent

Scan consumes axis 0, so x must have layout (T_total, D). Transposing it changes which dimension is treated as time. The hidden vector must retain its shape and dtype at every step. Returning a scalar in its place violates that contract; Python scalar outputs are not inherently forbidden when their array shape and dtype match the expected structure.

Build the cell once before scanning. Constructing it in the body is unnecessary and does not inherently create a fresh random seed each iteration. The explicit carry, not mutation of the cell object, stores this sequence's evolving hidden state.