--- jupytext: formats: md:myst text_representation: extension: .md format_name: myst format_version: 0.13 jupytext_version: 1.16.4 kernelspec: display_name: Python 3 language: python name: python3 --- (jax-201-debugging)= # 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 page introduces JAX's built-in debugging tools, which work with the various JAX transformations. **Summary:** - Use {func}`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 {func}`jax.debug.print` for traced (dynamic) array values with {func}`jax.jit`, {func}`jax.vmap` and others. - Use Python {func}`print` for static values, such as dtypes and array shapes. Recall from {doc}`jit` (and {ref}`jax-101-tracing`) that when transforming a function with {func}`jax.jit`, the Python code is executed with abstract tracers in place of your arrays. Because of this, Python's {func}`print` function prints only the tracer: ```{code-cell} 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.) ``` Python's `print` executes at trace time, before the runtime values exist. If you want to print the actual runtime values, you can use {func}`jax.debug.print`: ```{code-cell} @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.) ``` Similarly, within {func}`jax.vmap`, using Python's `print` will only print the tracer; to print the values being mapped over, use {func}`jax.debug.print`: ```{code-cell} 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) ``` Here's the result with {func}`jax.lax.map`, which is a sequential map rather than a vectorization: ```{code-cell} result = jax.lax.map(f, xs) ``` Notice the order is different, as {func}`jax.vmap` and {func}`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 {func}`jax.grad`, where {func}`jax.debug.print` only prints the forward pass. In this case, the behavior is similar to Python's {func}`print`, but unlike `print`, it stays the same if you also apply {func}`jax.jit`. ```{code-cell} def f(x): jax.debug.print("jax.debug.print(x) -> {}", x) return x ** 2 result = jax.grad(f)(1.) ``` When the printed values don't depend on one another, calls to {func}`jax.debug.print` may run in a different order once staged out by a JAX transformation. If you need the original order, such as `x: ...` first and then `y: ...` second, add the `ordered=True` parameter. For example: ```{code-cell} @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) ``` (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. For pure code that makes no difference; printing makes it visible.) ### Sharp bits of `jax.debug.print` A few cautions: **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 {func}`jax.custom_vjp`: ```{code-cell} @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.) ``` **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 problems: adding a print can change the values you're trying to inspect. Printing sharded values likewise forces synchronization to gather the value. **Prints are asynchronous.** Like the computations they're embedded in ({ref}`jax-201-async-dispatch`), debug prints can arrive after the enclosing function has returned, even after `block_until_ready()`, which waits for values but not for side effects. To wait for outstanding prints, use {func}`jax.effects_barrier`: ```{code-cell} @jax.jit def f(x): jax.debug.print("x: {}", x) return x f(2.).block_until_ready() jax.effects_barrier() ``` ## `jax.debug.callback` for more control during debugging {func}`jax.debug.print` is implemented using the more flexible {func}`jax.debug.callback`, which gives greater control over the host-side logic executed via a Python callback. It is compatible with {func}`jax.jit`, {func}`jax.vmap`, {func}`jax.grad` and other transformations (refer to the {ref}`jax-201-callbacks-flavors` table in {doc}`callbacks` for more information). For example: ```{code-cell} 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); ``` It works under {func}`jax.vmap` and {func}`jax.grad` too: ```{code-cell} x = jnp.arange(5.0) jax.vmap(f)(x); ``` ```{code-cell} jax.grad(f)(1.0); ``` This can make {func}`jax.debug.callback` useful for general-purpose debugging. You can learn more about {func}`jax.debug.callback` and other kinds of JAX callbacks in {doc}`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, makes computations raise an error as soon as a NaN is produced. 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, JAX re-runs the function in de-optimized op-by-op mode, effectively removing one level of `@jax.jit` at a time. Tricky situations can 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, but your code will continue to execute. If the NaNs are being produced in the backward pass of a gradient evaluation, then when an exception is raised, several frames up the stack trace you'll find JAX's backward pass (a function named like `backward_pass`), 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: ```{code-cell} 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)) ``` The NaN was caught as 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`: ```{code-cell} :tags: [raises-exception] jax.jit(f)(5.) ``` 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: ```{code-cell} with jax.debug_nans(False): print(jax.jit(f)(5.)) ``` ```{code-cell} 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 are expensive, so 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: ```python import jax import jax.numpy as jnp 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 ({ref}`jax-501-multiprocess`). 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 {doc}`slow-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 {doc}`callbacks`.