Skip to content

← Flax NNX step 5 of 100

Medium Primitives

NNX Mutability Demo

Jump to the task

Why this matters

A model can keep state between calls, such as a counter used for logging. Here, an NNX Variable stores a scalar that starts at zero and increases once per call. Reusing the same model lets the counter accumulate; creating a new model resets it.

Worked example

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

class Counter(nnx.Module):
    def __init__(self, rngs):
        _ = rngs.params()
        self.count = nnx.Variable(jnp.array(0.0, dtype=jnp.float32))

    def __call__(self, x):
        self.count[...] = self.count[...] + 1.0
        return x

model = Counter(rngs=nnx.Rngs(0))
x = jnp.array([1.0, 2.0, 3.0], dtype=jnp.float32)
assert jnp.array_equal(model(x), x)
assert jnp.array_equal(model(x), x)
assert model.count[...] == 2.0

The constructor consumes one params key to follow this exercise's setup. The counter does not use randomness, and consuming a key has no effect on its shape or value. An independent counter implementation needs no RNG.

Updating a Variable

variable[...] reads its array value. Assigning to variable[...] stores a replacement array in the existing wrapper. The JAX array itself remains immutable; the Variable is the mutable container.

Array immutability does not mean every use of += is invalid. Python can implement augmented assignment by computing a new value and assigning it back. This lesson uses explicit read-and-write syntax to make that clear. Avoid replacing the wrapper itself with a raw array: that changes how type filters see the value.

A counter records the event where you increment it. Incrementing inside a forward call counts calls, not optimizer updates. A Variable label does not by itself prevent differentiation; gradient and optimizer filters determine which values they operate on.

Execution and tracing

The exercise uses an ordinary Python loop. Reading an eager scalar after its update does not race with that update. Converting a traced scalar to a Python number is a different issue: it prevents JIT tracing.

NNX transformations such as nnx.jit support model-state updates. Plain jax.jit should not be used to mutate a model captured from an outer scope. Keep state explicit or use the NNX transformation appropriate to the task. Compiling a tiny counter is not a guaranteed speedup; compilation has a cost.

Problem

Implement mutability_increment_count(seed, x, n_calls):

  1. Define Counter as shown above. It increments the counter and returns the input unchanged.
  2. Create one instance using nnx.Rngs(int(seed)).
  3. Set last_x = x, then call the model on x int(n_calls) times, assigning each returned value to last_x.
  4. Return jnp.stack([model.count[...], last_x.sum()]).

Seed and n_calls are nonnegative integer-valued scalars; both may arrive as floats. The tested counts are small and exactly representable in float32. Input x is a finite float32 vector. Return a float32 array of shape (2,), containing [n_calls, x.sum()]. Initializing last_x before the loop makes zero calls valid: the result is [0.0, x.sum()].

Common pitfalls

  • Build the model once before the loop. Constructing it inside the loop resets the count each time.
  • Read the final counter value from the model. It is defined even when the loop makes no calls.
  • Keep dynamic values as JAX arrays when building the result. Do not convert the counter or an output sum to a Python float.
  • If JIT-compiling this setup function, mark seed and n_calls static. The Python loop is unrolled during tracing, so this is suitable only for small fixed counts, not an efficient dynamic training loop.

The numerical fixtures verify the returned values, not the use of a mutable module. Run the example's repeated calls to inspect the state changes while practicing the intended implementation.