Skip to content

← Flax NNX step 17 of 100

Medium Primitives

NNX Counters and Buffers

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 StepCounter(nnx.Module):
    def __init__(self, rngs):
        _ = rngs.params()
        self.step = nnx.Variable(jnp.array(0.0, dtype=jnp.float32))

    def __call__(self):
        self.step[...] = self.step[...] + 1.0
        return self.step[...]

model = StepCounter(rngs=nnx.Rngs(0))
assert model() == 1.0
assert model() == 2.0
assert model.step[...] == 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 counter_after_n_calls(seed, n_calls):

  1. Define StepCounter as shown above.
  2. Create one instance using nnx.Rngs(int(seed)).
  3. Call it int(n_calls) times in a Python loop.
  4. Return jnp.stack([model.step[...], jnp.float32(int(n_calls))]).

Seed and n_calls are nonnegative integer-valued scalars; both may arrive as floats. The tested call counts are small and exactly representable in float32. Return shape (2,), containing [final_step, n_calls]. With zero calls, both values are zero.

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.