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:
Understand why Python's assert silently fails inside jax.jit
Enforce shape and type contracts on your functions
Hunt down NaN values before they destroy a training run
Catch runtime value errors using chex.chexify
Detect expensive, silent recompilation before it wastes your compute budget
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:
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:
@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:
@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:
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
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:
# 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
# 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:
# 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:
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:
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:
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.
You don't have to wait until the gradient step. You can check activations inside your model's forward pass too:
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:
@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.
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:
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.
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:
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.
# 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.
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:
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:
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:
Static checks everywhere, always: assert_shape, assert_type, assert_rank on function boundaries. Cheap, catches most bugs, safe to leave in production.
assert_tree_all_finite at gradient computation: Place it right after value_and_grad, before the optimizer update. This is your NaN firewall.
chexify during development only: Use it to catch value-based logic errors while building a new component. Strip it before long production runs.
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.
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
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.
Add a trace budget: Decorate your train_step with @chex.assert_max_traces(n=2). Run training normally and confirm it passes.
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.
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
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
