# Catching bugs before they catch you

Here's a scenario every ML engineer eventually lives through: you kick off a training run before leaving the office. Three hours later, you check back in. The loss curve looks like a flat line at `NaN`. Somewhere in hour one, a single bad gradient poisoned every parameter in your model, and you didn't find out until the GPU had burned through three hours of compute for nothing.

This is the dark side of JAX's speed.

`@jax.jit` compiles your training loop into a single optimized kernel that runs on the GPU with almost no Python overhead. That's exactly what makes it fast and exactly what makes it dangerous. When something goes wrong inside a compiled function, JAX doesn't pause to tell you. It either crashes with a wall of cryptic XLA output, or it keeps running and quietly corrupts your model.

Today we fix that.

We're introducing **Chex**, a library built by DeepMind specifically for testing and validating JAX code. Chex lets you write "contracts" for your functions; assertions that check shapes, types, and numerical validity, and that actually work inside compiled JAX code, unlike standard Python `assert`.

By the end of this article, you'll be able to:

1.  Understand why Python's `assert` silently fails inside `jax.jit`
    
2.  Enforce shape and type contracts on your functions
    
3.  Hunt down NaN values before they destroy a training run
    
4.  Catch runtime value errors using `chex.chexify`
    
5.  Detect expensive, silent recompilation before it wastes your compute budget
    
6.  Build a self-validating neural network layer
    

Let's become site inspectors.

## Why standard Python assertions fail

In ordinary Python, NumPy, or PyTorch code, an assertion is simple:

```python
def process_batch(images):
    assert images.shape == (32, 28, 28, 1), "Wrong image shape!"
    return images * 2.0
```

This works because when the function runs, `images` is a real array with real values. Python can check `.shape` immediately and raise an error if it doesn't match.

Now wrap this in `@jax.jit`:

```python
@jax.jit
def process_batch(images):
    assert images.shape == (32, 28, 28, 1), "Wrong image shape!"
    return images * 2.0
```

Here's the problem.

When JAX JIT-compiles a function, it doesn't run it with real data the first time. It runs it with **abstract tracers**; placeholder objects that only know the *shape* and *dtype* of the eventual input, not the actual values.

Shape assertions like the one above *can* work during tracing, because shape is known statically. But the moment you try to assert something about the actual *values*, like checking whether a tensor contains NaNs, things fall apart:

```python
@jax.jit
def broken_check(x):
    # This will crash with a ConcretizationTypeError
    assert jnp.all(x > 0), "Values must be positive!"
    return x * 2
```

JAX will throw a `ConcretizationTypeError`, essentially saying: "I can't evaluate `x > 0` as a Python boolean right now, because `x` doesn't have real values yet, it's just a tracer."

This is the core problem Chex solves. It gives you assertion tools that understand the difference between the tracing phase and the execution phase, so your checks work reliably in both.

## Shape and type contracts

The single most common source of silent bugs in deep learning is the **silent broadcast**.

Imagine you have model predictions of shape `(32,)` and labels of shape `(32, 1)`. If you subtract them:

```python
predictions = jnp.zeros((32,))
labels = jnp.zeros((32, 1))

diff = predictions - labels  # No error!
print(diff.shape)  # (32, 32) completely wrong, but silent
```

NumPy and JAX both allow this "broadcasting" behavior. Instead of throwing an error, they expand the shapes to be compatible, quietly producing a `(32, 32)` matrix from what should have been a `(32,)` vector. Your loss function will average this garbage, and your model will train on nonsense; with no error message anywhere.

Chex fixes this by letting you enforce exact shape contracts.

### assert\_shape

```python
import jax.numpy as jnp
import chex

x = jnp.ones((32, 12))

# Exact shape check
chex.assert_shape(x, (32, 12))

# Use None for dimensions that can vary (e.g., batch size)
chex.assert_shape(x, (None, 12))

# This will raise an AssertionError with a clear message
chex.assert_shape(x, (32, 10))
```

### assert\_rank

Sometimes you don't care about exact dimensions, just how many there are:

```python
# Ensure x is a 2D matrix (batch, features)
chex.assert_rank(x, 2)

# Ensure x is either 2D or 3D
chex.assert_rank(x, (2, 3))
```

### assert\_type

```python
# Ensure the array is float32, not int32 or float64
chex.assert_type(x, jnp.float32)
```

This catches a subtle but common bug: accidentally passing integer labels where the model expects floats, which can cause silent precision issues or gradient problems.

### assert\_axis\_dimension

Sometimes you need to check just one specific axis:

```python
# Ensure axis 1 has exactly 12 elements (e.g., 12 financial features)
chex.assert_axis_dimension(x, axis=1, expected=12)
```

### Putting it together

Let's write a data validation function for a mobile money fraud detection pipeline:

```python
import jax
import jax.numpy as jnp
import chex

@jax.jit
def process_transactions(features, user_ids):
    # Type checks
    chex.assert_type(features, jnp.float32)
    chex.assert_type(user_ids, jnp.int32)
    
    # Rank checks
    chex.assert_rank(features, 2)   # (batch, num_features)
    chex.assert_rank(user_ids, 1)   # (batch,)
    
    # Shape checks: batch dimension is flexible, feature count is fixed
    chex.assert_shape(features, (None, 12))
    
    # Ensure features and user_ids have matching batch sizes
    chex.assert_equal_shape_prefix([features, user_ids], prefix_len=1)
    
    return features * 0.1  # Placeholder processing
```

If someone accidentally passes a `(64, 13)` feature matrix instead of `(64, 12)`, Chex catches it immediately during tracing, with a clear message pointing at exactly which dimension is wrong. No silent broadcasting, no three-hour debugging session later.

There's a second benefit here that's easy to overlook: these assertions are **documentation**. A teammate reading this function doesn't have to guess what shape `features` should be: the code tells them, and enforces it.

## Hunting NaNs

Shape bugs are annoying. NaN bugs are catastrophic.

A `NaN` (Not a Number) is viral. If one gradient in a 20-million-parameter model becomes `NaN`,maybe from a division by zero, a `log(0)`, or an exploding gradient, that NaN spreads through the optimizer state on the very next update. Within a step or two, every parameter in your model is `NaN`, and there's no way to recover without rolling back to an earlier checkpoint.

The hard part isn't knowing *that* your model died. It's knowing *when* and *why*.

### assert\_tree\_all\_finite

Because our models are built with Flax NNX, gradients live inside deeply nested pytrees; dictionaries of dictionaries containing every layer's weights. Chex gives us a way to check the entire tree in one call:

```python
import chex

# grads is a pytree: could be nested dicts, lists, NNX State objects, etc.
chex.assert_tree_all_finite(grads)
```

This single line recursively walks through every leaf array in the tree and checks that none of them contain `NaN` or `Inf`. If even one value out of millions is bad, it raises an error immediately.

Let's wire this into our training step from [week 4](https://kambale.dev/optax-optimizers):

```python
from flax import nnx
import optax
import chex

@nnx.jit
def bulletproof_train_step(model, optimizer, x, y):
    
    def loss_fn(m):
        logits = m(x)
        loss = optax.softmax_cross_entropy_with_integer_labels(logits, y).mean()
        return loss
    
    loss, grads = nnx.value_and_grad(loss_fn)(model)
    
    # The Firewall: check before we let bad gradients touch the model
    chex.assert_tree_all_finite(grads)
    
    optimizer.update(model, grads)
    return loss
```

By placing the check immediately after computing gradients and before applying them, we build a firewall. If a batch produces exploding gradients, the program crashes right then and there, with the exact batch of data still available for inspection; rather than three hours later when you notice the loss plot is flat.

### Checking intermediate activations

You don't have to wait until the gradient step. You can check activations inside your model's forward pass too:

```python
class SafeCNN(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)
        
        # Catch problems right where they happen
        chex.assert_tree_all_finite(x)
        
        x = x.reshape(x.shape[0], -1)
        return self.linear(x)
```

This is more expensive than checking only the final gradients (since it runs on every forward pass), but during development it can pinpoint exactly which layer is misbehaving.

## Runtime value checks with chexify

Shape, type, and finiteness checks all work during JAX's abstract tracing phase because they only need to know about the *structure* of your data, not its actual values.

But what if your check genuinely depends on runtime values? For example, verifying that a custom softmax implementation never produces negative probabilities:

```python
@jax.jit
def custom_softmax(x):
    probs = jax.nn.softmax(x)
    
    # This crashes with ConcretizationTypeError
    assert jnp.all(probs >= 0), "Probabilities cannot be negative!"
    
    return probs
```

As we saw earlier, this fails because `probs >= 0` can't be evaluated as a Python boolean during tracing—there are no concrete values yet.

`@chex.chexify` solves this by rewriting your function so that value-based assertions are deferred and evaluated safely at actual runtime, rather than during compilation.

```python
import chex

# chexify goes OUTSIDE jit

@chex.chexify
@jax.jit
def safe_softmax(x):
    probs = jax.nn.softmax(x)
    
    chex.assert_equal(
        jnp.all(probs >= 0), 
        True, 
        custom_message="Negative probability detected!"
    )
    
    return probs
```

There's a subtlety in how you call a chexified function: the assertion is checked asynchronously, so you need to explicitly wait for it to complete:

```python
result = safe_softmax(jnp.array([1.0, 2.0, -1.0]))
chex.block_until_chexify_assertions_complete()
```

### A word of caution

`chexify` is not free. It adds real overhead by inserting extra runtime checks into what would otherwise be pure, optimized XLA code. It's also worth noting it currently doesn't work in Colab notebooks; something to keep in mind if you're prototyping there.

The right mental model: `chexify` is scaffolding, not a permanent structure. Use it liberally during development and debugging. Strip it out (or gate it behind a debug flag) before a long production training run where every millisecond of compute matters.

```python
DEBUG_MODE = True

def maybe_chexify(fn):
    return chex.chexify(fn) if DEBUG_MODE else fn

@maybe_chexify
@jax.jit
def train_step(...):
    ...
```

## The silent assassin: recompilation

So far we've covered code that crashes loudly or corrupts silently. There's a third failure mode that's sneakier: code that runs, produces correct results, but runs *far* slower than it should.

Recall from Week 1 that `@jax.jit` compiles a specialized version of your function for each unique combination of input shapes and dtypes. The first time you call it with a `(32, 28, 28, 1)` batch, XLA compiles a kernel for that exact shape. If you then call it with a `(16, 28, 28, 1)` batch, say, because it's the last, incomplete batch of an epoch, JAX has to compile an entirely new kernel from scratch.

Compilation isn't cheap. It can take anywhere from milliseconds to several seconds depending on the complexity of your model. If your data pipeline is silently feeding inconsistent shapes, or worse, passing plain Python lists instead of JAX arrays, which forces recompilation based on the *values* of the list, your "fast, compiled" training loop can secretly spend most of its time recompiling instead of training.

This is invisible unless you're looking for it. Your loss curve looks fine. Your code doesn't crash. It's just mysteriously, consistently slow.

### assert\_max\_traces

Chex gives you a tripwire for this exact problem:

```python
import chex

@chex.assert_max_traces(n=2)
@nnx.jit
def strict_train_step(model, optimizer, batch):
    # ... training logic ...
    return loss
```

This decorator tracks how many times the function has triggered XLA compilation. If it's compiled more times than `n` allows, Chex raises an error immediately, telling you exactly what's happening: your function is recompiling more than expected.

A typical budget: `n=1` for a function that should compile exactly once, or `n=2` if you expect one compilation for a "warmup" call plus one for the real shape. Anything beyond that usually signals a bug in your data pipeline; most commonly, batches of inconsistent size.

The fix is almost always the same: make sure your data loader produces consistently shaped batches, typically by dropping the last incomplete batch (`drop_remainder=True` in TensorFlow Datasets, or the equivalent in Grain, which we'll cover in Week 8) or by padding it to match.

```python
# Clear the trace counter between tests if needed
chex.clear_trace_counter()
```

## The bulletproof NNX layer

Let's combine everything into a real, production-style component: a self-validating layer that defends itself against bad input at every stage.

This is the pattern you'll find in serious production ML systems; layers that don't just compute math, but actively verify their own contracts.

```python
from flax import nnx
import chex
import jax
import jax.numpy as jnp

class BulletproofBlock(nnx.Module):
    """A linear + ReLU block that validates its own inputs and outputs."""
    
    def __init__(self, in_features: int, out_features: int, rngs: nnx.Rngs):
        self.in_features = in_features
        self.out_features = out_features
        self.linear = nnx.Linear(in_features, out_features, rngs=rngs)
    
    def __call__(self, x):
        # Input contract
        chex.assert_type(x, jnp.float32)
        chex.assert_rank(x, 2)
        chex.assert_axis_dimension(x, axis=1, expected=self.in_features)
        
        # The actual math
        x = self.linear(x)
        x = nnx.relu(x)
        
        # Output contract
        chex.assert_rank(x, 2)
        chex.assert_axis_dimension(x, axis=1, expected=self.out_features)
        chex.assert_tree_all_finite(x)
        
        return x
```

Let's test it:

```python
rngs = nnx.Rngs(jax.random.PRNGKey(0))
safe_layer = BulletproofBlock(in_features=64, out_features=128, rngs=rngs)

# Correct input: should pass
good_data = jnp.ones((32, 64))
output = safe_layer(good_data)
print(f"Success! Output shape: {output.shape}")

# Wrong input: should fail with a clear error
bad_data = jnp.ones((32, 32))  # Wrong feature dimension

try:
    output = safe_layer(bad_data)
except Exception as e:
    print(f"Caught the expected error:\n{e}")
```

Expected output:

```plaintext
Success! Output shape: (32, 128)
Caught the expected error:
[Chex] Assertion assert_axis_dimension failed: expected axis 1 of shape (32, 32) to be 64
```

Notice what this buys you. If you're building something like a medical imaging classifier, where a shape mismatch or silent NaN could mean a wrong diagnosis slipping through unnoticed, you don't want to *hope* your tensors line up. You want the architecture itself to refuse to run on bad data.

## A practical debugging workflow

Here's how to think about applying these tools in practice, roughly in order of how expensive they are:

1.  **Static checks everywhere, always**: `assert_shape`, `assert_type`, `assert_rank` on function boundaries. Cheap, catches most bugs, safe to leave in production.
    
2.  `assert_tree_all_finite` **at gradient computation**: Place it right after `value_and_grad`, before the optimizer update. This is your NaN firewall.
    
3.  `chexify` **during development only**: Use it to catch value-based logic errors while building a new component. Strip it before long production runs.
    
4.  `assert_max_traces` **on your core training step**: Set it once you know your expected trace count, and let it catch data pipeline regressions automatically.
    
5.  `nnx.display()` **for structural sanity**: Not a Chex tool, but pairs well—use it to visually confirm your model's shape assumptions match what your Chex assertions expect.
    

## Exercises

1.  **Retrofit Week 4's training loop**: Add `chex.assert_tree_all_finite(grads)` right after your `nnx.value_and_grad` call in the OCR training step from Week 4.
    
2.  **Add a trace budget**: Decorate your `train_step` with `@chex.assert_max_traces(n=2)`. Run training normally and confirm it passes.
    
3.  **Sabotage your own pipeline**: Intentionally force your data loader to yield a batch of size 16 on the final iteration of an epoch (instead of your usual 32). Watch `assert_max_traces` catch the resulting recompilation. Then fix it by dropping the incomplete batch.
    
4.  **Build a NaN generator**: Write a tiny function that deliberately divides by zero inside a `loss_fn`. Confirm that `assert_tree_all_finite` catches the resulting NaN gradient before it reaches the optimizer.
    

## Quick reference

```python
import chex
import jax.numpy as jnp

# Shape and type contracts
chex.assert_shape(x, (None, 12))          # None = any size allowed
chex.assert_rank(x, 2)                     # Exactly 2 dimensions
chex.assert_type(x, jnp.float32)           # Enforce dtype
chex.assert_axis_dimension(x, axis=1, expected=12)
chex.assert_equal_shape_prefix([a, b], prefix_len=1)  # Matching batch dims

# NaN/Inf Hunting
chex.assert_tree_all_finite(pytree)        # Works on nested structures

# Runtime value assertions
@chex.chexify
@jax.jit
def fn(x):
    chex.assert_equal(jnp.all(x >= 0), True)
    return x

chex.block_until_chexify_assertions_complete()

# Recompilation detection
@chex.assert_max_traces(n=2)
@jax.jit
def fn(x):
    return x

chex.clear_trace_counter()  # Reset between tests

# Comparing PyTrees
chex.assert_trees_all_close(tree1, tree2, rtol=1e-5)
chex.assert_trees_all_equal(tree1, tree2)
```

## What's next

Chex catches the bugs that announce themselves loudly: wrong shapes, wrong types, exploding gradients. But some bugs are quieter; your model trains, the loss goes down, but the predictions are subtly wrong. No shape error, no NaN, just flawed logic hiding somewhere in your forward pass.

Next week, we go deeper into the compiled black box. We'll learn how to use `jax.debug.print` and `jax.debug.breakpoint` to inspect values *inside* JIT-compiled functions, how to temporarily disable JIT for standard `pdb` debugging, and how to profile your code to find performance bottlenecks that Chex can't catch.

**Next week**: *When JIT hides your errors: Debugging JAX*