Skip to content

← Flax NNX step 97 of 100

Medium Primitives

NNX Eager Debug

Jump to the task

Why this matters

Running a model without an outer JIT transformation makes ordinary Python inspection convenient. You can pause in the forward method, inspect its attributes and print intermediate arrays while investigating a bug.

This exercise gives a model an integer counter. Each forward call updates the counter and returns the output of a Linear layer. The counter changes, but the layer's parameters and input stay the same, so its output repeats.

Worked example

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

class CountingModel(nnx.Module):
    def __init__(self, in_features, features, rngs):
        self.l1 = nnx.Linear(in_features, features, rngs=rngs)
        self.counter = nnx.Variable(jnp.array(0, dtype=jnp.int32))

    def __call__(self, x):
        self.counter[...] = self.counter[...] + 1
        return self.l1(x)

model = CountingModel(3, 4, rngs=nnx.Rngs(0))
x = jnp.array([1.0, 2.0, 3.0], dtype=jnp.float32)
first = model(x)
second = model(x)
assert model.counter[...] == 2
assert jnp.array_equal(first, second)
print(int(model.counter[...]))  # 2

The counter persists because both calls use the same model instance. Assigning to the Variable's value preserves its wrapper while storing a replacement array. The underlying JAX array is still immutable.

The counter records forward calls, not optimizer steps. A parameter filter excludes this generic Variable; no optimizer runs in this exercise.

What changes under JIT?

A debugger or Python print inside a compiled function generally runs when Python traces that function, not once for every execution of cached code. Use JAX runtime debugging tools such as jax.debug.print when you need values from compiled execution. Eager and compiled debugging are both useful.

NNX can carry Variable updates through nnx.jit. Continuing the example:

@nnx.jit
def apply_model(model, inputs):
    return model(inputs)

apply_model(model, x)
apply_model(model, x)
assert model.counter[...] == 4

The state update happens on both executions. Do not interpret Python tracing as an extra logical forward call. Keep the counter as a JAX value within the compiled function; converting it to a Python int or float there needs a concrete value and can fail during tracing.

A plain Python integer attribute can survive split/merge as static graph information. It is not the dynamic array state used here. Conversely, using a Variable does not make arbitrary Python side effects work under JIT.

Problem

Implement eager_debug_step_count(seed, x, features, n_calls):

  1. Define CountingModel as shown above, constructing its Linear layer with default initializer and bias settings.
  2. Build it with input width x.shape[-1], output width int(features) and nnx.Rngs(int(seed)).
  3. Call the same model on x int(n_calls) times in an ordinary Python loop, retaining the last output.
  4. Return a float32 array containing the final counter and the sum of the last output. Use JAX stacking rather than Python float conversions.

Input x is a nonempty finite float32 vector. Features and n_calls are positive integer-valued scalars; seed is a nonnegative integer-valued scalar. These scalars may arrive as floats. At least one call is required because the function returns information from the last call. Return shape (2,): [n_calls, last_output.sum()].

Common pitfalls

  • Constructing a fresh model inside the loop resets the counter.
  • Replacing the Variable attribute with a raw array changes its wrapper and filter behavior. Update counter[...] explicitly.
  • The counter must change on every call even though the Linear output repeats. Repeated numerical outputs do not imply unchanged model state.
  • JIT is not a guaranteed speedup for every workload. Validate behavior and measure performance before choosing where to compile.
  • For optional JIT testing of the whole setup function, make seed, features and n_calls static. Its Python loop is then unrolled at tracing.

Numerical fixtures cannot establish how often a submitted model was called. Run the example's state assertions to inspect the intended behavior.