jax.lax.scan inside a Flax Module

Repeat a transformation with one shared matrix

Define a Linen module that owns one learned matrix W, then use jax.lax.scan to repeat a pure array computation with that matrix. The scan body reads W but does not create or update Linen variables.

This is a useful separation: Linen manages the parameter, while JAX manages the loop. When a loop needs to call stateful Linen modules, a lifted transformation such as nn.scan handles their variables and random-number streams. Here, the body only needs matrix multiplication and tanh.

Count the initial transform as well as the loop

For x of shape (H,) and W of shape (H, H), first compute:

carry_0 = tanh(x @ W)

Then run exactly T additional updates:

carry_(k+1) = tanh(carry_k @ W), for k = 0, ..., T-1

Return carry_T. The complete function applies the matrix-and-tanh transformation T + 1 times. In particular, T=0 returns tanh(x @ W), not x. No new input vector arrives at each step.

Separate carry from collected outputs

A scan body returns two things: the carry for the next step and an optional output for the current step. In this exercise:

body(carry, unused) -> (tanh(carry @ W), None)

The first value advances the recurrence. Returning None as the second value means there is no output history to stack. Consequently, the second result of lax.scan is also None; the first result is the final (H,) carry.

Use jnp.arange(T) as the scan input. Its leading length sets the number of iterations, but the body ignores its values. The alternative xs=None, length=T expresses the same loop; follow the stated form for this exercise.

A recurrence with known weights

This example runs two scan steps after the initial transformation and checks the answer with explicit array operations:

import jax
import jax.numpy as jnp

x = jnp.array([.5, -.25])
W = jnp.array([[1., .5], [0., 1.]])
initial = jnp.tanh(x @ W)

def body(carry, unused):
    return jnp.tanh(carry @ W), None

final, history = jax.lax.scan(body, initial, jnp.arange(2))
expected = jnp.tanh(jnp.tanh(initial @ W) @ W)
assert jnp.allclose(final, expected)
assert final.shape == (2,)
assert history is None

zero_step, _ = jax.lax.scan(body, initial, jnp.arange(0))
assert jnp.array_equal(zero_step, initial)

Shared does not mean frozen

Declare W before defining or calling the scan body. Every iteration uses that same array, so the module has H * H scalar parameters regardless of T. Capturing W in the body's closure does not stop gradients: differentiation through the module accounts for its use in the initial transform and in every scan step.

Keep Linen variable creation outside this raw JAX transformation. Do not call self.param or create a parameterized submodule inside the body. Plain array operations avoid mixing Linen scope side effects with JAX tracing; the exact error from an invalid mixture can depend on how it is written.

The carry must retain its pytree structure, shape and dtype at every step. Here it remains one float vector of length H. These requirements apply to scan regardless of whether a lifted wrapper manages parameters.