Skip to content

← JAX Conceptual Deep-Dives step 12 of 15

Medium Primitives

jax.debug.callback

Jump to the task

Python tracing and runtime callbacks are different

A Python function body runs when JAX traces it. Its ordinary print statements therefore do not report new dynamic values on every cached execution. jax.debug.callback instead stages a call that can inspect concrete values when the compiled calculation runs.

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

def show_mean(arr):
    mean = np.asarray(arr).mean()
    print(f"mean of squared values: {mean:.4f}")

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

for values in ([1., 2., 3.], [2., 3., 4.]):
    compute(jnp.array(values)).block_until_ready()
    jax.effects_barrier()
# mean of squared values: 4.6667
# mean of squared values: 9.6667

Both inputs have the same shape and dtype, so the second call can reuse the compiled function while its callback receives different values. The mean is calculated on the host for inspection. It is not returned into the compiled sum.

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 debug_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.