Debugging runtime values

Debugging runtime values#

Do you have exploding gradients? Are NaNs making you gnash your teeth? Just want to poke around the intermediate values in your computation? This section introduces you to a set of built-in JAX debugging methods that you can use with various JAX transformations.

Summary:

Use jax.debug.print() to print values to stdout in jax.jit-, jax.vmap-, and other transformation-decorated functions.
JAX offers config flags and context managers that enable catching errors more easily. For example, enable the jax_debug_nans flag to automatically detect when NaNs are produced in jax.jit-compiled code and enable the jax_disable_jit flag to disable JIT-compilation.

`jax.debug.print` for simple inspection#

Here is a rule of thumb:

Use jax.debug.print() for traced (dynamic) array values with jax.jit(), jax.vmap() and others.
Use Python print() for static values, such as dtypes and array shapes.

Recall from Just-in-time compilation (and How transformations work: tracing) that when transforming a function with jax.jit(), the Python code is executed with abstract tracers in place of your arrays. Because of this, the Python print() function will only print this tracer value:

import jax
import jax.numpy as jnp

@jax.jit
def f(x):
  print("print(x) ->", x)
  y = jnp.sin(x)
  print("print(y) ->", y)
  return y

result = f(2.)

print(x) -> JitTracer(~float32[])
print(y) -> JitTracer(~float32[])

Python’s print executes at trace-time, before the runtime values exist. If you want to print the actual runtime values, you can use jax.debug.print():

@jax.jit
def f(x):
  jax.debug.print("jax.debug.print(x) -> {x}", x=x)
  y = jnp.sin(x)
  jax.debug.print("jax.debug.print(y) -> {y}", y=y)
  return y

result = f(2.)

jax.debug.print(x) -> 2.0
jax.debug.print(y) -> 0.9092974066734314

Similarly, within jax.vmap(), using Python’s print will only print the tracer; to print the values being mapped over, use jax.debug.print():

def f(x):
  jax.debug.print("jax.debug.print(x) -> {}", x)
  y = jnp.sin(x)
  jax.debug.print("jax.debug.print(y) -> {}", y)
  return y

xs = jnp.arange(3.)

result = jax.vmap(f)(xs)

jax.debug.print(x) -> 0.0
jax.debug.print(x) -> 1.0
jax.debug.print(x) -> 2.0
jax.debug.print(y) -> 0.0
jax.debug.print(y) -> 0.8414709568023682
jax.debug.print(y) -> 0.9092974066734314

Here’s the result with jax.lax.map(), which is a sequential map rather than a vectorization:

result = jax.lax.map(f, xs)

jax.debug.print(y) -> 0.0
jax.debug.print(x) -> 0.0
jax.debug.print(y) -> 0.8414709568023682
jax.debug.print(x) -> 1.0
jax.debug.print(y) -> 0.9092974066734314
jax.debug.print(x) -> 2.0

Notice the order is different, as jax.vmap() and jax.lax.map() compute the same results in different ways. When debugging, the evaluation order details are exactly what you may need to inspect.

Below is an example with jax.grad(), where jax.debug.print() only prints the forward pass. In this case, the behavior is similar to Python’s print(), but it’s consistent if you apply jax.jit() during the call.

def f(x):
  jax.debug.print("jax.debug.print(x) -> {}", x)
  return x ** 2

result = jax.grad(f)(1.)

jax.debug.print(x) -> 1.0

Sometimes, when the arguments don’t depend on one another, calls to jax.debug.print() may print them in a different order when staged out with a JAX transformation. If you need the original order, such as x: ... first and then y: ... second, add the ordered=True parameter.

For example:

@jax.jit
def f(x, y):
  jax.debug.print("jax.debug.print(x) -> {}", x, ordered=True)
  jax.debug.print("jax.debug.print(y) -> {}", y, ordered=True)
  return x + y

f(1, 2)

jax.debug.print(x) -> 1
jax.debug.print(y) -> 2

Array(3, dtype=int32, weak_type=True)

(The reordering happens because the compiler receives a functional representation of the staged-out computation, in which the imperative order of your Python statements is gone and only data dependence remains — invisible for pure code, visible once printing enters the picture.)

Sharp bits of `jax.debug.print`#

A few cautions worth knowing before you sprinkle jax.debug.print everywhere:

Format strings are deferred. The format string can’t be an f-string: f-strings format immediately, while jax.debug.print needs to delay formatting until the runtime value exists. Pass values as arguments, as in the examples above.

Printing on the backward pass takes an extra step. As shown above, jax.debug.print fires on the forward pass only. To see gradients, wrap a print in a jax.custom_vjp():

@jax.custom_vjp
def print_grad(x):
  return x

def print_grad_fwd(x):
  return x, None

def print_grad_bwd(_, x_grad):
  jax.debug.print("x_grad: {}", x_grad)
  return (x_grad,)

print_grad.defvjp(print_grad_fwd, print_grad_bwd)

def f(x):
  x = print_grad(x)
  return x * 2.

jax.grad(f)(1.)

x_grad: 2.0

Array(2., dtype=float32, weak_type=True)

Debug prints perturb the computation. Adding jax.debug.print changes the program XLA compiles: a value that would have lived inside a fused kernel must be materialized so it can be printed, which can change performance, memory usage, and even numerics (different fusions can round differently). Keep this in mind when debugging numerical mysteries — the act of looking can disturb the thing you’re looking at. Printing sharded values likewise forces synchronization to gather the value.

Prints are asynchronous. Like the computations they’re embedded in (Asynchronous dispatch), debug prints can arrive after the enclosing function has returned — even after block_until_ready(), which waits for values, not side effects. To wait for outstanding prints, use jax.effects_barrier():

@jax.jit
def f(x):
  jax.debug.print("x: {}", x)
  return x

f(2.).block_until_ready()
jax.effects_barrier()

x: 2.0

`jax.debug.callback` for more control during debugging#

jax.debug.print() is implemented using the more flexible jax.debug.callback(), which gives greater control over the host-side logic executed via a Python callback. It is compatible with jax.jit(), jax.vmap(), jax.grad() and other transformations (refer to the Flavors of callback table in External callbacks for more information).

For example:

import logging

def log_value(x):
  logging.warning(f'Logged value: {x}')

@jax.jit
def f(x):
  jax.debug.callback(log_value, x)
  return x

f(1.0);

WARNING:root:Logged value: 1.0

This callback is compatible with other transformations, including jax.vmap() and jax.grad():

x = jnp.arange(5.0)
jax.vmap(f)(x);

WARNING:root:Logged value: 0.0
WARNING:root:Logged value: 1.0
WARNING:root:Logged value: 2.0
WARNING:root:Logged value: 3.0
WARNING:root:Logged value: 4.0

jax.grad(f)(1.0);

WARNING:root:Logged value: 1.0

This can make jax.debug.callback() useful for general-purpose debugging.

You can learn more about jax.debug.callback() and other kinds of JAX callbacks in External callbacks.

Throwing Python errors with JAX’s debug flags#

JAX offers flags and context managers that enable catching errors more easily: jax_debug_nans to automatically detect when NaNs are produced in jax.jit-compiled code, and jax_disable_jit to disable JIT-compilation, enabling the use of traditional Python debugging tools like print and pdb.

`jax_debug_nans`#

jax_debug_nans is a JAX flag that when enabled, will cause computations to error-out immediately on production of a NaN. Switching this option on adds a NaN check to every floating point type value produced by XLA. That means values are pulled back to the host and checked as ndarrays for every primitive operation not under an @jax.jit.

For code under an @jax.jit, the output of every @jax.jit function is checked and if a NaN is present it will re-run the function in de-optimized op-by-op mode, effectively removing one level of @jax.jit at a time.

There could be tricky situations that arise, like NaNs that only occur under a @jax.jit but don’t get produced in de-optimized mode. In that case you’ll see a warning message print out but your code will continue to execute.

If the NaNs are being produced in the backward pass of a gradient evaluation, when an exception is raised several frames up in the stack trace you will be in the backward_pass function, which is essentially a simple jaxpr interpreter that walks the sequence of primitive operations in reverse.

To turn on the NaN-checker, do one of:

run your code inside the jax.debug_nans context manager, using with jax.debug_nans(True):;
set the JAX_DEBUG_NANS=True environment variable;
add jax.config.update("jax_debug_nans", True) near the top of your main file;
add jax.config.parse_flags_with_absl() to your main file, then set the option using a command-line flag like --jax_debug_nans=True.

For example:

import traceback
jax.config.update("jax_debug_nans", True)

def f(x):
  w = 3 * jnp.square(x)
  return jnp.log(-w)

# The stack trace is very long, so only print a couple of lines.
try:
  f(5.)
except FloatingPointError as e:
  print(traceback.format_exc(limit=2))

Invalid nan value encountered in the output of a jax.jit function. Calling the de-optimized version.
Traceback (most recent call last):
  File "/tmp/ipykernel_1589/1275719523.py", line 10, in <module>
    f(5.)
  File "/tmp/ipykernel_1589/1275719523.py", line 6, in f
    return jnp.log(-w)
           ^^^^^^^^^^^
FloatingPointError: invalid value (nan) encountered in log

The NaN generated was caught, with an ordinary Python exception — so running %debug in IPython gives a post-mortem debugger at the fault. This also works with functions under @jax.jit:

jax.jit(f)(5.)

Invalid nan value encountered in the output of a jax.jit function. Calling the de-optimized version.
Invalid nan value encountered in the output of a jax.jit function. Calling the de-optimized version.

---------------------------------------------------------------------------
FloatingPointError                        Traceback (most recent call last)
Cell In[13], line 1
----> 1 jax.jit(f)(5.)

    [... skipping hidden 5 frame]

Cell In[12], line 6, in f(x)
      4 def f(x):
      5   w = 3 * jnp.square(x)
----> 6   return jnp.log(-w)

    [... skipping hidden 5 frame]

File ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/numpy/ufuncs.py:491, in log(x)
    456 @export
    457 @jit(inline=True)
    458 def log(x: ArrayLike, /) -> Array:
    459   """Calculate element-wise natural logarithm of the input.
    460 
    461   JAX implementation of :obj:`numpy.log`.
   (...)    489     Array(True, dtype=bool)
    490   """
--> 491   out = lax.log(*promote_args_inexact('log', x))
    492   jnp_error._set_error_if_nan(out)
    493   return out

    [... skipping hidden 7 frame]

File ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/pjit.py:176, in _run_python_pjit(p, args_flat, fun, args, kwargs)
    174 except api_util.InternalFloatingPointError as e:
    175   if getattr(fun, '_apply_primitive', False):
--> 176     raise FloatingPointError(
    177         f"invalid value ({e.ty}) encountered in {fun.__qualname__}") from None
    178   api_util.maybe_recursive_nan_check(e, fun, args, kwargs)  # should always raise.
    179   raise RuntimeError("Internal error") from e  # fall-back error to be safe.

FloatingPointError: invalid value (nan) encountered in log

When a NaN appears in the output of an @jax.jit function, JAX re-runs the de-optimized code, so we still get a clear stack trace pointing at the producing operation. The jax.debug_nans context manager can scope the checking; since we activated it globally above, let’s deactivate it:

with jax.debug_nans(False):
  print(jax.jit(f)(5.))

nan

jax.config.update("jax_debug_nans", False)  # back off for the rest of this page

Strengths: easy to apply; precisely detects where NaNs were produced; throws a standard Python exception and is compatible with PDB postmortem.

Limitations: re-running functions eagerly can be slow, and the constant device-to-host checks cost real performance — don’t leave the NaN-checker on when you’re not debugging. It also errors on false positives, e.g. intentionally created NaNs.

`jax_debug_infs`#

jax_debug_infs works similarly to jax_debug_nans. It often needs to be combined with jax_disable_jit, since Infs might not cascade to the output the way NaNs do.

`jax_disable_jit`#

jax_disable_jit is a JAX flag that when enabled, disables JIT-compilation throughout JAX (including in control flow functions like jax.lax.cond and jax.lax.scan). With compilation out of the picture, your function is plain Python running eagerly, so all the ordinary tools work: print, pdb, Python’s built-in breakpoint().

You can disable JIT-compilation by:

running your code inside the jax.disable_jit context manager, using with jax.disable_jit():;
setting the JAX_DISABLE_JIT=True environment variable;
adding jax.config.update("jax_disable_jit", True) near the top of your main file;
adding jax.config.parse_flags_with_absl() to your main file, then setting the option using a command-line flag like --jax_disable_jit=True.

For example:

import jax
jax.config.update("jax_disable_jit", True)

def f(x):
  y = jnp.log(x)
  if jnp.isnan(y):
    breakpoint()
  return y

jax.jit(f)(-2.)  # ==> Enters PDB breakpoint!

Strengths: easy to apply; enables Python’s built-in breakpoint and print; throws standard Python exceptions and is compatible with PDB postmortem.

Limitations: running functions without JIT-compilation can be slow.

Warning

These flags are best suited to single-process development, and don’t work well in multi-controller (multi-process) JAX (Introduction to multi-controller JAX (aka multi-process/multi-host JAX)). Raising a Python error on one process but not the others — say, when only one process’s shard produces a NaN under jax_debug_nans — breaks the assumption that every process runs the same program in lockstep, and the usual symptom is the remaining processes hanging in a collective.

Next steps#

Check out Debugging slow JAX tracing and XLA compilation if the thing that needs debugging is tracing or compile time itself. For the mechanism underlying jax.debug.print and jax.debug.callback, see External callbacks.

Debugging runtime values

Contents

Debugging runtime values#

jax.debug.print for simple inspection#

Sharp bits of jax.debug.print#

jax.debug.callback for more control during debugging#

Throwing Python errors with JAX’s debug flags#

jax_debug_nans#

jax_debug_infs#

jax_disable_jit#

Next steps#

`jax.debug.print` for simple inspection#

Sharp bits of `jax.debug.print`#

`jax.debug.callback` for more control during debugging#

`jax_debug_nans`#

`jax_debug_infs`#

`jax_disable_jit`#