Tracing and Debugging JAX Functions

1. The Cache Key: What Triggers Recompilation?

As we discussed, the first time you call a jit-decorated function, JAX performs tracing and compilation, which is slow. It then caches the resulting machine code. Subsequent calls are fast because they hit this cache. The critical question is: what constitutes the "key" for this cache? If a function call generates a key that JAX hasn't seen before, it will trigger a re-trace and re-compilation.

The diagram below illustrates this. The "FIRST CALL (Slow Path)" is what we want to minimize. The "FAST PATH! (cache hit)" is where the performance gains are.

JAX determines whether to recompile by inspecting the signature of the function's inputs. A new compilation is triggered if you call the function with a different:

1. Number of array arguments.
2. Shape of an array argument.
3. dtype of an array argument.
4. Pytree structure of the arguments (e.g., passing (x, y) instead of [x, y]).

Let's read a quick summary that lists these triggers.

The most common trigger in practice is a change in array shape. For example, in deep learning, the final batch of a dataset often has a different size if you don't use drop_last=True in your data loader. Calling a jitted training step with this smaller batch will trigger a new compilation.

2. Diagnosing the Trace with jax.make_jaxpr

To effectively debug tracing and compilation issues, you need to see what JAX sees. The result of the tracing process is a Jaxpr (JAX Program Representation). This is a simple, functional intermediate representation (IR) that captures the sequence of primitive operations your function performs. Given your background, you can think of it as conceptually similar to other compiler IRs you may have encountered, or to the graph representation produced by torch.jit.trace.

JAX provides a utility, jax.make_jaxpr, that lets you generate and inspect the jaxpr for a function without actually running the full compilation and execution.

import jax
import jax.numpy as jnp

def my_func(x, y):
  a = jnp.sin(x)
  b = jnp.cos(y)
  return a + b




# Create example inputs to define the shape/dtype for tracing
x_example = jnp.ones(3)
y_example = jnp.zeros(3)




# Generate and print the jaxpr
jaxpr = jax.make_jaxpr(my_func)(x_example, y_example)
print(jaxpr)

This will output:

{ lambda ; a:f32[3] b:f32[3]. let
    c:f32[3] = sin a
    d:f32[3] = cos b
    e:f32[3] = add c d
  in (e,) }

Let's break this down:

• lambda ; a:f32[3] b:f32[3].: This defines the function signature. The inputs x and y are represented by abstract variables a and b with their respective types (f32[3]).
• let ...: This block defines the body of the computation. Each line is a primitive operation (sin, cos, add) that binds its result to a new variable (c, d, e).
• in (e,): This specifies the return value(s) of the function, in this case, the result of the addition.

This jaxpr is the pure, static computation graph that JAX hands off to the XLA compiler. Notice how all Python-level constructs are gone; only the essential JAX primitives remain.

3. Controlling Tracing with static_argnums

In the last lesson, you saw that Python control flow depending on a traced value, like if x > 0:, fails because x is a Tracer with no concrete value at trace-time. But what if you genuinely need to use an argument's value to define the structure of the computation graph?

Consider a function that creates an array of a given length:

def create_filled(val, length):



  # We need the *value* of `length` here
  return jnp.full((length,), val)




# This will fail!
# jitted_create_filled = jax.jit(create_filled)
# jitted_create_filled(jnp.pi, 5) 

This fails because when JAX traces the function, length becomes a Tracer, but the shape argument to jnp.full must be a concrete Python integer, not a tracer.

The solution is static_argnums. It's an argument to jax.jit that lets you specify which positional arguments should not be traced. Instead, their concrete Python values are kept and are available during the trace.

As the article demonstrates, by using jit_create_filled = jax.jit(create_filled, static_argnums=(1,)), you are telling JAX: "When you trace this function, the first argument (val) should be a Tracer, but the second argument (length, at index 1) should be treated as a static, compile-time constant."

When you inspect the jaxpr for jit_create_filled(2, 5), you see that 5 is baked directly into the graph: b:i32[5] = .... When you inspect it for jit_create_filled(1.6, 10), 10 is baked in: b:f32[10] = ....

This brings us to the crucial trade-off:

• Pro: static_argnums lets you use Python control flow (if, for, list comprehensions, etc.) that depends on the value of an argument, because that value is now concrete at trace-time.
• Con: You have added the static argument's value to the cache key. Every time you call the function with a new value for that static argument, you trigger a recompilation.

This is only a good strategy if the number of possible values for the static argument is small and fixed. A boolean flag is a perfect use case, as it only ever results in two compiled versions of the function. A hyperparameter like num_layers in a model is also a good candidate if you only experiment with a few values. Using it for an argument that changes frequently, like batch size, can destroy performance due to constant recompilation.

Let's watch a short video segment that reinforces these ideas with more examples, including how Python for loops are "unrolled" during tracing when their length depends on a static argument.

Conclusion

You now have a much more sophisticated understanding of JAX's compilation model. You can not only explain what tracing is but also predict when it will be re-triggered and how to control it.

Key Takeaways:

• Function recompilation is triggered by changes to the input signature: array shapes, dtypes, or pytree structures.
• You can diagnose what JAX will compile by using jax.make_jaxpr to inspect the jaxpr, which is the pure computation graph generated during the trace.
• You can force an argument to remain a concrete Python value during tracing by using static_argnums. This is necessary for using that argument's value in Python control flow.
• Using static_argnums adds the argument's value to the compilation cache key. This is powerful but can lead to excessive recompilation if the static argument has many possible values.

In the next lesson, we will introduce JAX's third major transformation: jax.vmap. This function is JAX's idiomatic solution for vectorization and is fundamental to efficiently processing batches of data. We will then explore how to compose jit, grad, and vmap together, which is where the true expressive power of JAX's programming model begins to shine.