flags.md

JAX debugging flags

JAX offers flags and context managers that enable catching errors more easily.

jax_debug_nans configuration option and context manager

Summary: Enable the jax_debug_nans flag to automatically detect when NaNs are produced in jax.jit-compiled code (but not in jax.pmap or jax.pjit-compiled code).

jax_debug_nans is a JAX flag that when enabled, automatically raises an error when a NaN is detected. It has special handling for JIT-compiled -- when a NaN output is detected from a JIT-ted function, the function is re-run eagerly (i.e. without compilation) and will throw an error at the specific primitive that produced the NaN.

Usage

If you want to trace where NaNs are occurring in your functions or gradients, you can turn on the NaN-checker by:

• setting the JAX_DEBUG_NANS=True environment variable;
• adding jax.config.update("jax_debug_nans", True) near the top of your main file;
• adding jax.config.parse_flags_with_absl() to your main file, then set the option using a command-line flag like --jax_debug_nans=True;

Example(s)

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

def f(x, y):
  return x / y
jax.jit(f)(0., 0.)  # ==> raises FloatingPointError exception!

Strengths and limitations of jax_debug_nans

Strengths

• Easy to apply
• Precisely detects where NaNs were produced
• Throws a standard Python exception and is compatible with PDB postmortem

Limitations

• Not compatible with jax.pmap or jax.pjit
• Re-running functions eagerly can be slow
• Errors on false positives (e.g. intentionally created NaNs)

jax_disable_jit configuration option and context manager

Summary: Enable the jax_disable_jit flag to disable JIT-compilation, enabling use of traditional Python debugging tools like print and pdb

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

Usage

You can disable JIT-compilation by:

• 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 set the option using a command-line flag like --jax_disable_jit=True;

Examples

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 and limitations of jax_disable_jit

Strengths

• Easy to apply
• Enables use of Python's built-in breakpoint and print
• Throws standard Python exceptions and is compatible with PDB postmortem

Limitations

• Not compatible with jax.pmap or jax.pjit
• Running functions without JIT-compilation can be slow