Skip to content

← Flax NNX step 96 of 100

Medium Primitives

NNX Debug Callback

Jump to the task

Why this matters

jax.debug.callback passes runtime values to a Python function on the host. It is useful when a custom diagnostic needs more than a formatted print. The callback should return None; it does not supply a result to the array computation.

Observe a runtime value

import jax
import jax.numpy as jnp

observed = []

def record_norm(value):
    observed.append(float(value))

@jax.jit
def double_with_diagnostic(x):
    jax.debug.callback(record_norm, jnp.linalg.norm(x))
    return 2 * x

y = double_with_diagnostic(jnp.array([3.0, 4.0], dtype=jnp.float32))
y.block_until_ready()
jax.effects_barrier()
assert jnp.allclose(y, jnp.array([6.0, 8.0]))
print(observed)  # [5.0] in this simple, single-call example

The callback receives a concrete value, so converting this scalar to a Python float is appropriate. Wait for effects before inspecting the list. This local debugging example is not a dependable metric-storage system.

What a debug callback does not guarantee

Transformations may drop, duplicate or reorder debug callbacks. Setting ordered=True orders supported debug effects relative to one another; it does not make dispatch synchronous or guarantee exactly one record per step. Avoid using debug callbacks for required file writes or other essential I/O.

For required impure host operations, investigate jax.experimental.io_callback and its transformation restrictions. For a pure host computation whose result belongs in the calculation, use jax.pure_callback with its result shape and dtype specification. Neither is an interchangeable replacement for debugging.

Callback data may need transfer to the host. Keep diagnostics small and remove them when measuring performance. A debug callback can affect execution even when its Python body does nothing.

Problem

Implement debug_callback_passthrough(seed, x, features):

  1. Define an NNX model with two nnx.Linear layers. The first maps x.shape[-1] to int(features); the second maps int(features) to itself.
  2. Create one nnx.Rngs(int(seed)) and pass it to both constructors, creating l1 before l2. Keep default initializers and biases.
  3. In __call__, compute h = self.l1(x), call jax.debug.callback(lambda v: None, jnp.linalg.norm(h)), and return self.l2(h). Do not add an activation function.
  4. Construct the model and return model(x).

The callback intentionally does nothing with the norm. It does not change the specified arithmetic, though debug instrumentation can affect compilation and floating-point details. Compare numerical results with a tolerance.

This entrypoint runs eagerly; the callback API also supports use inside nnx.jit. Numerical tests check the returned array and cannot establish that the callback ran. Include it as an implementation requirement.

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),).