Skip to content

← Flax NNX step 95 of 100

Medium Primitives

NNX Debug Print

Jump to the task

Why this matters

Python print inside a traced function usually shows a tracer rather than the input's numerical values. Use jax.debug.print to inspect values during execution. In an eager call, ordinary Python printing can still be useful.

A complete example

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

class DebugModel(nnx.Module):
    def __init__(self, in_features, features, rngs):
        self.l1 = nnx.Linear(in_features, features, rngs=rngs)
        self.l2 = nnx.Linear(features, features, rngs=rngs)

    def __call__(self, x):
        h = self.l1(x)
        jax.debug.print("layer 1 norm: {n}", n=jnp.linalg.norm(h))
        return self.l2(h)

model = DebugModel(3, 4, rngs=nnx.Rngs(0))

@nnx.jit
def forward(model, x):
    return model(x)

x = jnp.array([1.0, 2.0, 3.0], dtype=jnp.float32)
y = forward(model, x)
y.block_until_ready()
jax.effects_barrier()
assert y.shape == (4,)

The diagnostic is the Euclidean norm of the first layer's output. The returned array is the second layer's output. No activation function belongs between these layers in this exercise.

Formatting and execution

Pass a format string and its values separately, as above. An f-string would format its contents too early, during tracing. jax.debug.print returns None, so keep the model's return statement.

Independent debug prints may be reordered. ordered=True preserves relative order among ordered debug effects where supported; it does not make the call synchronous, and ordered effects are unsupported under pmap. Debug effects can be duplicated or dropped by transformations. They are for inspection, not a reliable count of training steps.

Printing may finish after the array result is ready. jax.effects_barrier() waits for outstanding effects. Debugging can also change timing, memory use and compilation optimizations, so remove diagnostics before benchmarking.

Problem

Implement debug_print_passthrough(seed, x, features):

  1. Create one nnx.Rngs(int(seed)). Pass it to both layer constructors, creating l1 before l2. Keep default initializers and biases.
  2. Define an NNX model with l1 mapping x.shape[-1] to int(features) and l2 mapping int(features) to int(features).
  3. In __call__, compute h = self.l1(x), print its norm using jax.debug.print("layer 1 norm: {n}", n=jnp.linalg.norm(h)), and return self.l2(h).
  4. Construct the model and return model(x).

The entrypoint calls the model eagerly. The example above shows how the same model works with nnx.jit. Numerical tests check the returned array, not the printed message; include the diagnostic even though it does not set the result.

Inputs: seed is a nonnegative integer-valued scalar supplied as a float; features is a positive integer-valued scalar supplied as a float. Cast both to int. x is a nonempty, finite float32 vector.

Output: a vector of shape (int(features),).