Topolearnprint and pdb do not see runtime values under jit, so you need JAX's own debugging tools.
Inside jit, a normal print(x) shows a tracer, not values. JAX offers runtime-aware tools:
jax.debug.print("x = {x}", x=x)prints actual values each time the compiled function runs. It works underjit,vmap,gradandscan. Note the format string uses named or positional braces, not an f-string.jax.debug.breakpoint()drops into a debugger with live values.jax.debug.callback(fn, *args)calls arbitrary Python with concrete arrays, e.g. to log to a file.
For numeric problems:
jax.config.update("jax_debug_nans", True)makes any operation producing a NaN raise immediately; JAX then re-runs the function un-jitted to point at the offending line. It slows things down, so enable it only while hunting.jax.config.update("jax_disable_jit", True)(or thejax.disable_jit()context manager) runs everything eagerly so ordinaryprintandpdbwork.jax.experimental.checkifyadds functional error checks (checkify.check(x > 0, "msg")) that survivejit.
Two classic NaN sources: jnp.where with an unsafe branch (the gradient of the unused branch is still computed, so jnp.where(x > 0, jnp.log(x), 0.) gives NaN gradients at x <= 0; apply the where to the input first), and jnp.sqrt at exactly zero in a norm.
Ordering of jax.debug.print output is not guaranteed across independent operations; pass ordered=True if order matters.
Example
import jax
import jax.numpy as jnp
@jax.jit
def f(x):
y = jnp.sin(x) * 2
jax.debug.print("y = {y}", y=y) # prints real values at run time
return y.sum()
f(jnp.array([0.0, jnp.pi / 2])) # y = [0. 2.]
def bad_log(x):
return jnp.where(x > 0, jnp.log(x), 0.0)
print(jax.grad(bad_log)(0.0)) # nan: log's gradient is still evaluated
def safe_log(x):
safe_x = jnp.where(x > 0, x, 1.0) # sanitise the input first
return jnp.where(x > 0, jnp.log(safe_x), 0.0)
print(jax.grad(safe_log)(0.0)) # 0.0
with jax.disable_jit():
print(jax.jit(lambda x: x + 1)(1.0)) # runs eagerly: 2.0Remember
jax.debug.printandjax.debug.breakpointsee runtime values under jitjax_debug_nanspinpoints the first NaN;disable_jitmakes ordinary tools work- Guard the inputs of
jnp.wherebranches to avoid NaN gradients
Watch out for
- Using an f-string with jax.debug.print, which formats the tracer at trace time instead of the runtime value.