Skip to content

← JAX Profiling & Debugging step 4 of 12

Medium Primitives

jax.debug.print with Multiple kwargs

Jump to the task

Print related values together

A sum alone does not tell you the typical size of an array's entries. Logging its mean on the same line gives that context. One debug.print call can format both values, avoiding two separate callback calls and keeping their labels together.

import jax
import jax.numpy as jnp

@jax.jit
def inspect_values(x):
    total = jnp.sum(x)
    mean = jnp.mean(x)
    jax.debug.print("sum={s} mean={m}", s=total, m=mean)
    return total

result = inspect_values(jnp.array([1., 2., 3.]))
result.block_until_ready()
jax.effects_barrier()
# stdout: sum=6.0 mean=2.0
# result is the scalar JAX array 6.0

The Python variable names do not have to match the placeholders: total is passed under the keyword s, so {s} prints total. The keyword names supplied to debug.print must match the format.

Check the format before chasing numerical bugs

Under normal format checking, a missing named value raises KeyError, and an unused formatting keyword raises ValueError. Extra values are not silently ignored. These checks can fail immediately in eager code or during JIT tracing, before the runtime callback executes.

Keep the format string plain and pass dynamic values separately. An f-string formats too early when traced. You can use format specifications, such as {m:.2f}, when a rounded scalar display helps; display rounding does not change the returned sum.

Debug output normally goes to stdout and may be captured by the surrounding runtime. It can arrive asynchronously, even after the result array is ready. The host effects_barrier above waits for pending effects before inspecting the output.

Combining values into one print keeps that message together; it does not impose a global ordering across separate calls or devices. Debug callbacks have overhead and can affect compilation, so use them for diagnosis rather than timing or guaranteed application logging.

Problem

Implement debug_print_two_things(x):

  1. Compute s = jnp.sum(x) and m = jnp.mean(x).
  2. Call jax.debug.print("sum={s} mean={m}", s=s, m=m).
  3. Return s, not m or the print call's None result.

X is a one-dimensional float32 JAX array. Return a scalar JAX array. For [10, 20], print sum=30.0 mean=15.0 and return 30. For an empty array, the sum is zero and the mean is NaN: print those values and return zero, without inventing a zero mean.

Output tests check the returned sum. They do not verify that both values were printed, so inspect the format and keyword arguments too.