Skip to content

← JAX Profiling & Debugging step 5 of 12

Medium Primitives

jax.debug.callback for Side Effects

Jump to the task

Inspect a runtime value with Python

jax.debug.callback lets a Python function inspect concrete values from a JAX computation. It is useful for diagnostic output that needs more control than a format string. The callback should return None; it does not supply a numerical result to the surrounding JAX calculation.

import jax
import jax.numpy as jnp
import numpy as np

def show_norm(arr):
    norm = np.linalg.norm(np.asarray(arr))
    print(f"squared values have norm {norm:.4f}")

@jax.jit
def compute(x):
    y = x ** 2
    jax.debug.callback(show_norm, y)
    return jnp.sum(y)

result = compute(jnp.array([1., 2., 3.]))
result.block_until_ready()
jax.effects_barrier()
# squared values have norm 9.8995
# result is 14

The callback sees [1, 4, 9]. Its norm is sqrt(1 + 16 + 81), whereas the returned sum is 1 + 4 + 9. These are different quantities. The f-string is appropriate inside this host callback because its values are concrete. It is not formatting a tracer during JIT tracing.

Diagnostic effects are not required application work

Debug callbacks can be dropped, duplicated or reordered by JAX transformations. Do not use callback counts as authoritative training step counts, or rely on them for required writes or external actions. They can also run asynchronously. The host effects_barrier in the example waits for outstanding effects; waiting only for an array is not a general guarantee that its callback has finished.

pure_callback is for a pure host calculation that supplies values back to JAX. io_callback is for impure host work with stronger execution guarantees; its ordering is an explicit option, not automatic. Neither is an interchangeable replacement for every debug callback.

Gradients follow the numerical calculation

A debug callback can coexist with grad: observing x does not prevent JAX from differentiating a separate expression such as x ** 2. There is no returned callback value through which to propagate a derivative. The host callback's Python calculation is not itself differentiated, and observing a primal value does not automatically show its gradient.

Callback transfers and host work add overhead and can change compiler optimization. Use them for diagnosis, not for timing the surrounding calculation or imposing a global execution order.

Problem

Implement callback_passthrough(x):

  1. Define a callback cb(arr) whose body is pass.
  2. Call jax.debug.callback(cb, x).
  3. Return x + 1 as a JAX array.

X is a one-dimensional float32 JAX array. Return the same shape and dtype, adding one to each element. [1, 2, 3] becomes [2, 3, 4]; an empty array remains empty. Do not modify the input or return the callback's None result.

The no-op callback deliberately prints nothing. This exercise practices the API wiring; numerical output tests do not prove callback execution. Inspect the callback call as well. Keep any demonstration barrier in host code outside the submitted function.