When JIT hides your errors: Debugging JAX

Last week, Chex gave us a firewall against shape mismatches and NaN gradients. That firewall catches the bugs that announce themselves; a tensor with the wrong shape, a gradient that explodes into infinity.

But not every bug is loud. Sometimes your model trains without a single error. The loss goes down. The shapes all line up. And yet the predictions are subtly, mysteriously wrong. Somewhere in your forward pass, the logic itself is flawed; and there's no assertion for "this math is conceptually incorrect."

When that happens, your first instinct is probably the one every Python developer reaches for: sprinkle in a few print() statements, or drop a breakpoint() in the middle of the function. In standard Python, this works instantly.

Inside a @jax.jit-compiled function, it doesn't work the way you expect, and understanding exactly why is the key to debugging JAX effectively.

Today we open up the compiled black box. By the end of this article, you'll know how to:

1. Understand why standard debugging tools fail inside JIT

2. Print real runtime values from inside compiled functions with jax.debug.print

3. Pause execution mid-computation with jax.debug.breakpoint

4. Temporarily disable compilation entirely to use pdb like normal Python

5. Automatically catch the exact operation that produces a NaN

6. Use nnx.display() and Module.sow() to inspect Flax NNX models

7. Combine all of these into a systematic debugging workflow

Let's go hunting.

Why your debugger doesn't work inside JIT

To understand JAX debugging, you have to remember the two-phase life of a @jax.jit function, which we first touched on in week 1.

Phase 1: Tracing. The first time you call a JIT-compiled function, JAX runs it once; not with your real data, but with abstract tracers. A tracer knows the shape and dtype of an eventual array, but not its actual values. JAX uses this trace to build a computation graph.

Phase 2: Compilation and Execution. XLA takes that graph, compiles it into optimized machine code, and only then does your real data flow through it; as a single fused kernel, without Python in the loop at all.

Here's the problem for debugging: standard Python tools assume every line of code executes with real, inspectable values, one at a time, with Python in control the whole way through.

def broken_debug(x):
    print(f"x = {x}")       # Prints an abstract tracer, not a value
    if x > 0:                # Can crash, tracers don't have a boolean truth value
        return x * 2
    return x

jax.jit(broken_debug)(5.0)

Run this, and two things go wrong:

1. print(f"x = {x}") prints something like Traced<ShapedArray(float32[])>, not the number 5.0. During tracing, that's genuinely all x is.

2. if x > 0: may raise a ConcretizationTypeError. JAX can't evaluate a tracer as a Python boolean, because during tracing it doesn't know what the value will actually be at runtime.

pdb.set_trace() and breakpoint() suffer the same fate, they can show you a tracer's shape and dtype, but not the concrete numbers flowing through your model at runtime.

This isn't a bug in JAX. It's the direct cost of the speed we've been chasing since Week 1. Compilation requires JAX to reason about your function abstractly, once, before any real data exists. The tools we need have to be built for that reality.

jax.debug.print: The JAX-native print()

The most direct fix is jax.debug.print. It's designed to be embedded into the compiled computation graph itself, so it prints real, concrete values at actual runtime, even inside jit, vmap, and grad.

import jax
import jax.numpy as jnp

@jax.jit
def debug_forward(x, w):
    logits = jnp.dot(x, w)
    jax.debug.print("logits = {}", logits)
    return jax.nn.softmax(logits)

x = jnp.array([[1.0, 2.0, 3.0]])
w = jnp.array([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]])

output = debug_forward(x, w)

Output:

logits = [[2.2 2.8]]

That's a real, concrete array, not a tracer. The print statement got compiled into the XLA graph as an actual print operation, so it fires with real data every time the function runs.

Keeping output in order

Because JAX can reorder and parallelize operations for performance, prints inside heavily parallel code (like vmap) can come out of sequence. If you need strictly ordered output; for example, watching a value evolve across training steps, use ordered=True:

jax.debug.print("step {} loss = {}", step, loss, ordered=True)

Printing inside a training step

This is where jax.debug.print earns its keep; dropped directly into a real NNX training step:

from flax import nnx
import optax

@nnx.jit
def train_step(model, optimizer, batch):
    def loss_fn(model):
        logits = model(batch['image'])
        loss = optax.softmax_cross_entropy_with_integer_labels(
            logits=logits, labels=batch['label']
        ).mean()
        return loss, logits
    
    (loss, logits), grads = nnx.value_and_grad(loss_fn, has_aux=True)(model)
    
    # Peek at real numbers, right inside the compiled step
    jax.debug.print("loss = {loss}", loss=loss)
    jax.debug.print("max logit = {m}", m=jnp.max(logits))
    
    optimizer.update(model, grads)
    return loss

Every single call to train_step, compiled or not, will print the real loss and the real max logit. No recompilation needed, no disabling JIT, no performance penalty worth worrying about for occasional prints.

jax.debug.breakpoint: An interactive pause inside JIT

Sometimes a print statement isn't enough. You want to actually stop execution and poke around. jax.debug.breakpoint() gives you that, working like a lightweight version of pdb that's compatible with tracing.

@jax.jit
def suspicious_function(x, w):
    hidden = jnp.dot(x, w)
    
    jax.debug.breakpoint()  # Execution pauses here at RUNTIME
    
    return jax.nn.relu(hidden)

output = suspicious_function(x, w)

When this runs, you'll drop into an interactive (jaxdb) prompt:

Entering jaxdb:
(jaxdb) p hidden
Array([[2.2, 2.8]], dtype=float32)
(jaxdb) c

The command set is intentionally limited compared to full pdb:

Command	Effect
p <var>	Print a variable's real value
c	Continue execution
q	Quit

It's not a full debugger, but it's enough to answer the question that matters most: "What does this tensor actually look like right now?"

Making it conditional

You rarely want to break on every single call; usually you want to break only when something specific goes wrong, like a NaN appearing:

import jax.numpy as jnp
from jax import lax

@jax.jit
def conditional_breakpoint(x):
    has_nan = jnp.any(jnp.isnan(x))
    
    lax.cond(
        has_nan,
        lambda: jax.debug.breakpoint(),
        lambda: None,
    )
    
    return x * 2

Now the breakpoint only fires on the exact batch that introduces a NaN — instead of interrupting every single training step.

jax.disable_jit: Bringing back standard Python tools

Sometimes you don't want JAX's limited debugging vocabulary; you want your full, familiar Python debugger back: pdb, your IDE's breakpoints, real print() statements, everything.

jax.disable_jit() forces JAX into fully eager execution, exactly like plain NumPy. Every operation runs immediately with real values, so all your normal tools work again.

import jax

def my_function(x, w):
    hidden = jnp.dot(x, w)
    print(f"hidden = {hidden}")   # Works normally now!
    breakpoint()                   # Standard Python debugger works!
    return jax.nn.relu(hidden)

with jax.disable_jit():
    output = my_function(x, w)

You can also flip this globally, which is convenient while iterating in a notebook:

# Global flag
jax.config.update("jax_disable_jit", True)

# Or via environment variable, before starting Python
# JAX_DISABLE_JIT=1 python train.py

The trade-off

jax.disable_jit() gives you back full debugging power at the cost of everything that makes JAX fast; no fusion, no compilation, no XLA optimization. On a real training step, this can be 10-100x slower.

Treat it as a scalpel, not your default mode:

# Debugging: eager execution, full pdb access
with jax.disable_jit():
    debug_output = train_step(model, optimizer, tiny_batch)

# Production: full speed
output = train_step(model, optimizer, real_batch)

One limitation worth knowing: disable_jit() doesn't always let you step inside the internals of jax.vmap or jax.scan: those transformations have their own internal execution logic that isn't simply "turned off" by disabling JIT.

Automatic NaN detection

We spent last week manually hunting NaNs with chex.assert_tree_all_finite. JAX also has a built-in, lower-level tool for the same problem: the jax_debug_nans flag.

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

With this enabled, the moment any operation produces a NaN, JAX automatically re-runs that specific operation eagerly, outside of JIT, so it can point you directly at the exact line and operation responsible.

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

@jax.jit
def buggy_function(x):
    y = jnp.log(x)  # log(0) = -inf, log(negative) = NaN
    return y * 2

# This will halt immediately with a traceback pointing at jnp.log
result = buggy_function(jnp.array([-1.0, 2.0, 3.0]))

Instead of your model silently going to NaN somewhere deep inside a 50-layer network and you having no idea which layer did it, JAX interrupts execution the instant the NaN is born and shows you exactly where.

Limitations

• It significantly slows down executio; this is a debugging flag, not something to leave on in production.

• It doesn't work with jax.shard_map.

• It can occasionally flag NaNs that are actually intentional (some numerically stable algorithms briefly produce and then discard NaN-like intermediate values), so a small amount of judgment is still required.

Use it the same way you'd use chexify: turned on while chasing a bug, turned off before a long production run.

nnx.display: Seeing your model's structure

Shifting away from runtime values, sometimes the bug isn't in the math; it's in the architecture itself. Maybe a layer has the wrong shape, or a parameter didn't initialize the way you expected.

nnx.display() is the Flax NNX equivalent of print(pytorch_model), but considerably richer:

from flax import nnx

class CNN(nnx.Module):
    def __init__(self, *, rngs: nnx.Rngs):
        self.conv1 = nnx.Conv(1, 32, kernel_size=(3, 3), rngs=rngs)
        self.conv2 = nnx.Conv(32, 64, kernel_size=(3, 3), rngs=rngs)
        self.linear = nnx.Linear(3136, 10, rngs=rngs)
    
    def __call__(self, x):
        x = nnx.relu(self.conv1(x))
        x = nnx.relu(self.conv2(x))
        x = x.reshape(x.shape[0], -1)
        return self.linear(x)

model = CNN(rngs=nnx.Rngs(0))
nnx.display(model)

In a notebook, this renders as an interactive, expandable tree showing every submodule, every parameter's exact shape and dtype, and how everything nests together. It's the fastest way to confirm your architecture actually matches what you think you built — invaluable when a shape mismatch several layers deep is causing a downstream Chex assertion to fail, and you're not sure which layer is actually at fault.

Module.sow: Capturing intermediate activations

Sometimes you don't just want to print an intermediate value onc; you want to capture it for later inspection, without changing your model's function signature or return value.

Module.sow() lets a layer stash a value as an attribute on itself during the forward pass:

from flax import nnx

class InspectableCNN(nnx.Module):
    def __init__(self, *, rngs: nnx.Rngs):
        self.conv1 = nnx.Conv(1, 32, kernel_size=(3, 3), rngs=rngs)
        self.linear = nnx.Linear(21632, 10, rngs=rngs)
    
    def __call__(self, x):
        x = self.conv1(x)
        x = nnx.relu(x)
        
        # Stash this activation for later inspection
        self.sow(nnx.Intermediate, 'conv1_activation', x)
        
        x = x.reshape(x.shape[0], -1)
        return self.linear(x)

model = InspectableCNN(rngs=nnx.Rngs(0))
dummy_input = jnp.ones((4, 28, 28, 1))
output = model(dummy_input)

# Access what got "sown"
print(model.conv1_activation.value.shape)  # (4, 26, 26, 32)

This is especially useful for checking whether a specific activation is behaving reasonably, is it saturating at zero (dead ReLUs)? Is its scale exploding layer over layer?, without restructuring your __call__ method to return every intermediate value explicitly.

By default, calling the module multiple times appends to a tuple rather than overwriting, so you can also use sow to accumulate values across a sequence (useful inside RNN-style loops).

Putting it together: A debugging workflow

Here's a practical order of operations for a real bug, roughly cheapest-to-most-invasive:

1. Structural sanity check. Run nnx.display(model) first. Confirm the architecture is actually what you intended: shapes, layer order, parameter counts.

2. Static checks. Make sure last week's Chex assertions (assert_shape, assert_rank, assert_type) are in place at your function boundaries. Many "mysterious" bugs are actually silent broadcasts.

3. NaN detection. If the loss goes to NaN, turn on jax_debug_nans first — let JAX point you directly at the offending operation before you start manually inserting print statements.

4. Targeted prints. Once you have a rough idea of where the problem lives, add a few jax.debug.print calls around it to watch real values move through the computation.

5. Interactive breakpoint. If prints aren't enough, you need to actually poke at a value's shape or compare two tensors interactively, drop in jax.debug.breakpoint(), ideally gated with a lax.cond so it only fires on the problematic case.

6. Full eager debugging. If you're truly stuck and need your entire familiar toolchain, pdb, IDE breakpoints, stepping line by line, wrap the smallest possible piece of code in jax.disable_jit() and accept the performance hit temporarily.

7. Clean up before production. Strip out debug prints, disable jax_debug_nans, remove any disable_jit() context managers. None of these should ship in your real training run.

Exercises

1. The silent tracer: Write a JIT-compiled function with a plain print(x) inside it. Call it twice with different-shaped inputs. Notice that the print only fires during tracing (once per unique shape), not on every call; this is a common source of confusion. Then replace it with jax.debug.print and observe that it fires on every call.

2. The NaN hunt: Write a function that computes 1.0 / (x - 5.0) for a batch of values, one of which is exactly 5.0. Turn on jax_debug_nans and run it. Confirm JAX stops you at the exact division.

3. Conditional breakpoint: Take your Week 4 training step and add a lax.cond-gated jax.debug.breakpoint() that only triggers when the loss exceeds some threshold (say, 10.0). Feed it a batch designed to produce a large loss and confirm it pauses.

4. Sow and inspect: Add self.sow(...) calls after each convolutional block in your Week 3 CNN. After a forward pass, print the mean and standard deviation of each captured activation. Are any of them suspiciously close to zero (dead layer) or huge (exploding activations)?

Quick reference

import jax
import jax.numpy as jnp
from jax import lax

# Runtime Printing Inside JIT
jax.debug.print("value = {}", x)
jax.debug.print("step {} loss {}", step, loss, ordered=True)

# Interactive Breakpoint Inside JIT
jax.debug.breakpoint()

# Conditional breakpoint
lax.cond(has_nan, lambda: jax.debug.breakpoint(), lambda: None)

# Full Eager Execution (standard pdb/print work)
with jax.disable_jit():
    result = my_function(x)

jax.config.update("jax_disable_jit", True)  # Global toggle

# Automatic NaN Detection
jax.config.update("jax_debug_nans", True)

# Flax NNX Structural Inspection
nnx.display(model)

# Capturing Intermediate Values
self.sow(nnx.Intermediate, 'my_activation', x)
model.my_activation.value  # Access after forward pass

What's next

We can now debug both the shape of our bugs (Chex, week 6) and the runtime behavior of our bugs (this week). But there's a whole different category of problem we haven't touched: what happens when your model is correct, your training loop is correct, and everything is just... slow?

Next week, we tackle the unglamorous but critical problem of getting data into your model fast enough to keep the GPU fed. We'll introduce Grain, JAX's purpose-built data loading library, and build a high-performance pipeline that shuffles, augments, and batches data without ever letting your accelerator sit idle.

Next week: Feeding the beast: Data pipelines with grain