Skip to content

← JAX Conceptual Deep-Dives step 6 of 15

Easy Primitives

jax.debug.print

Jump to the task

Inspect a value inside compiled code

Python print runs when Python executes. Inside a jitted function, that usually means tracing: a dynamic array appears as a tracer, and a cached call does not execute that Python print again. Static information such as a shape can still be useful to print while tracing.

jax.debug.print lets you inspect runtime array values instead. Here both calls have the same shape and dtype, so the second can reuse the compilation while still printing its own values:

import jax
import jax.numpy as jnp

@jax.jit
def compute(x):
    y = x ** 2
    jax.debug.print("y = {y}", y=y)
    return jnp.sum(y)

compute(jnp.array([1., 2., 3.])).block_until_ready()
jax.effects_barrier()  # finish outstanding debug output
# y = [1. 4. 9.]
compute(jnp.array([4., 5., 6.])).block_until_ready()
jax.effects_barrier()
# y = [16. 25. 36.]

Formatting happens later

Pass a plain format string and the values separately: jax.debug.print("sum: {s}", s=s). The name s in the placeholder matches the keyword argument. Do not use an f-string: it formats immediately, before the dynamic value exists during tracing. Positional placeholders are supported too, but this exercise uses a named placeholder.

By default, debug.print writes to stdout, although a runtime or test runner may capture that stream. The call produces diagnostic output and returns None; return your computed array separately.

Debug output is not an execution guarantee

Debug callbacks can be asynchronous. Waiting on a returned array does not necessarily finish its prints; use jax.effects_barrier() from the host when you need outstanding effects to finish. Independent prints can be reordered, and transformations can duplicate or drop debug callbacks. Do not use them for required application side effects.

Prints add overhead and can change compiler optimization, memory use and even rounding. Remove them before comparing performance. They do not turn the whole calculation into synchronous execution.

Problem

Implement debug_printed_compute(x):

  1. Compute s = jnp.sum(x).
  2. Call jax.debug.print("sum: {s}", s=s).
  3. Return s as a scalar JAX array.

X is a one-dimensional float32 JAX array. For [1, 2, 3], print sum: 6.0 and return 6. For an empty input, print and return zero. Do not return the print call's None result or a formatted string.

The numerical tests check the sum, not captured output. Inspect the requested print call too. The effects_barrier calls above belong to the demonstration's host code, not inside your submitted function.