NNX Fori Loop in Module

Repeat a state update

This module starts with state=x and applies the same update T times:

state = tanh(state @ W + x)

W is one initialized matrix, shared across all iterations. x has two roles: it is the starting state and the fixed input added at every step. There is no sequence of different input vectors and no training step. The result is the final state, not a list of intermediate states.

jax.lax.fori_loop(0, T, body, x) expresses the loop. Its body receives an iteration index and the current state, then returns the next state. This recurrence does not use the index. The upper bound is exclusive, so nonnegative T gives exactly T updates; T=0 returns x unchanged.

Keep the loop body on arrays

Store W in nnx.Param, then read its array as W = self.W[...] before defining the body. The body uses that array and x while its state is passed explicitly through the loop. The state must keep the same shape and dtype after each update.

Reading W does not detach it from automatic differentiation. When the module is passed to an NNX gradient transformation, the array operations can still depend on its parameters. Extracting the array here makes the loop's inputs clear; it is not a claim that all read-only access to a parameter inside a loop is forbidden.

Do not mutate a captured module or RNG stream inside this raw JAX loop. If an algorithm needs changing array state, include it explicitly in the carry. This exercise updates only the activation, leaving W fixed.

A complete one-dimensional example

The example uses W=0.5 so the first two updates are easy to inspect. For the exercise, keep the specified random initialization instead.

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

class Recurrence(nnx.Module):
    def __init__(self):
        self.W = nnx.Param(jnp.array([[0.5]]))

    def __call__(self, x, T):
        W = self.W[...]
        def body(i, state):
            return jnp.tanh(state @ W + x)
        return jax.lax.fori_loop(0, T, body, x)

model = Recurrence()
x = jnp.array([0.2])
first = jnp.tanh(0.3)
second = jnp.tanh(0.5 * first + 0.2)
assert jnp.array_equal(model(x, 0), x)
assert jnp.allclose(model(x, 1), jnp.array([first]))
assert jnp.allclose(model(x, 2), jnp.array([second]))

Static and dynamic loop counts

A Python loop inside JIT can expand into repeated copies of its body. fori_loop can represent the repetition without fully unrolling it. With bounds known during tracing, JAX implements fori_loop using scan and supports reverse-mode differentiation. With dynamic bounds it uses a while loop, which does not support reverse-mode differentiation in this form. The choice depends on how the bounds reach the transform.

The exercise entrypoint converts T with int(T), so keep T static when applying JIT to that entrypoint. The module's loop can accept a dynamic integer bound when called separately without that Python conversion. Seed and hidden size also need to be static for the entrypoint.

Scan is another valid way to express a final-state recurrence: its body can return (next_state, None) to avoid collecting per-step outputs. Neither loop form guarantees less memory or faster execution for every workload. This exercise specifically practices fori_loop.

Bounded values do not guarantee convergence

After an update with finite inputs, tanh bounds each state component between -1 and 1. That does not guarantee the state stays near zero or converges to a fixed point. Scaling a random W by 0.1 does not ensure its matrix norm is below one, and a large x can saturate tanh. A matrix spectral norm below one is a sufficient contraction condition for this recurrence, but it is not a requirement checked by the task. Compute exactly T updates; do not stop early based on apparent convergence.