JAX Tracing: Understanding Tracers, Trace-time, and Run-time

1. The Protagonist of Compilation: The Tracer Object

When you call a function decorated with @jax.jit for the first time, a subtle but critical substitution occurs. JAX intercepts the input arrays you provide and, instead of passing them directly to your function, it replaces them with special placeholder objects called Tracers.

A Tracer is not an array with data; it is an abstract representation of one. It doesn't know the values contained within the array, but it knows everything else needed to define the computation: its shape and dtype. This process is a form of abstract interpretation. JAX executes your Python function not with concrete numbers, but with these abstract stand-ins to "trace" the sequence of operations.

As the documentation states, these tracer objects are used to extract the sequence of operations. Each jax.numpy function, when called on a Tracer, doesn't perform a calculation. Instead, it records the operation and outputs a new Tracer representing the result. For example, jnp.tanh(tracer_in) returns tracer_out, which has the same shape and dtype but simply has a "tanh" operation added to its history.

2. Trace-Time vs. Run-Time

This tracing mechanism cleanly separates the execution of your program into two distinct phases:

Trace-Time (or Compile-Time):

• What happens? The Python code of your function is executed once. JAX passes Tracers as inputs and records every jnp operation they encounter.
• What is the goal? To build a language-agnostic computation graph, called a jaxpr (JAX Expression). This graph captures the pure mathematical essence of your function.
• Side Effects: Python-native side effects, like print() statements or logging, will execute during this phase and only this phase.

Run-Time:

• What happens? The jaxpr generated during trace-time is compiled into highly optimized machine code by the XLA compiler. This compiled code is then executed on the accelerator (GPU/TPU) with the actual, concrete input data.
• What is the goal? To perform the numerical computation as fast as possible.
• Python's Role: The original Python code is no longer involved. The accelerator is executing a pre-compiled binary. This happens on the second and all subsequent calls to the function, as long as the input signature matches the cached version.

Patrick Kidger frames this beautifully: "the Python code is just specifying a way to build a computation graph." The trace-time phase builds the graph; the run-time phase executes it.

3. The Consequence: Python Control Flow

This separation has a critical implication that trips up many newcomers. Since trace-time logic operates on Tracers which are value-agnostic, any standard Python control flow that depends on the values inside an array will fail.

Consider this code:

@jax.jit
def f(x):
  if x > 0:  # This will fail!
    return x
  else:
    return -x

When JAX traces this function, x is a Tracer. The expression x > 0 is ambiguous. Is the tracer greater than zero? JAX cannot know, because the tracer represents any possible array of a given shape and dtype. It has no concrete value to check. JAX will raise a ConcretizationError to tell you that a tracer was used in a context that requires a concrete value.

The solution is to use JAX's own control flow primitives, which are themselves operations that get recorded in the graph. Instead of a Python if, you use jax.numpy.where:

@jax.jit
def f(x):



  # This works!
  return jnp.where(x > 0, x, -x)

Here, jnp.where is just another node in the computation graph. At trace-time, it's recorded as "a where operation will happen here." The actual branching logic is then handled efficiently by the compiled code at run-time, when the concrete value of x is known.

4. Advanced Topic: Static vs. Traced Operations

You can gain even finer control by understanding that not everything inside a jitted function has to be a traced value. Some values can be treated as static: their values are known at compile-time and are baked into the compiled program.

A prime example is an array's shape. Even if x is a Tracer, x.shape is a static Python tuple of integers. You can use these static values in Python control flow. However, if you convert a static value into a JAX array (e.g., jnp.array(x.shape)), it becomes a traced value, and the same rules apply.

This distinction gives rise to a common JAX idiom: using standard NumPy for operations you want to happen at trace-time on static values, and using jax.numpy for operations you want to be part of the compiled, run-time computation.

import numpy as np # For static, compile-time logic
import jax.numpy as jnp # For traced, run-time logic

@jax.jit
def get_center_pixel(image_matrix):



  # image_matrix is a Tracer
  # image_matrix.shape is a static tuple, e.g., (28, 28)
  



  # Static computation using standard Python/NumPy
  height, width = image_matrix.shape
  center_h = height // 2
  center_w = width // 2
  



  # The result is used in a traced operation
  return image_matrix[center_h, center_w] 

In this example, the calculation of the center indices is performed in Python during tracing. The final indexing operation image_matrix[...] is recorded in the jaxpr using the computed static indices.

Finally, it's worth noting that not all tracers are created equal. While jax.jit uses abstract tracers (ShapedArray) that are value-agnostic, other transformations like jax.grad use concrete tracers (ConcreteArray) that do carry value information along with shape and dtype. This is necessary for automatic differentiation, which often depends on the values of the inputs. We will explore this in depth in Module 4.

Conclusion

In this lesson, we've pulled back the curtain on JAX's compilation process. You now have a more rigorous mental model of how jax.jit works.

Key Takeaways:

• JAX execution is split into trace-time and run-time.
• Trace-time is a one-off Python execution with placeholder Tracer objects to build a computation graph (jaxpr).
• Run-time is the repeated execution of the compiled graph on an accelerator with real data.
• This model forbids standard Python control flow (if, for) from depending on the values of traced arrays. You must use jnp.where, lax.cond, etc., to embed logic into the graph itself.
• You can explicitly separate static, compile-time logic (using Python/NumPy) from traced, run-time logic (using jax.numpy).

We've established that JAX traces a function to create a compiled version. But what happens if we call the function with arguments that don't match the cached version? This triggers a re-trace and re-compilation, which can silently kill performance. In our next lesson, we will investigate the causes of re-compilation and learn how to diagnose and control this behavior using static_argnums and jax.make_jaxpr.