Tracing

Call a jitted function and JAX does not run your Python the way you might picture. It calls your function once, with tracers standing in for the real arguments: objects that carry shape and dtype but no actual values. Every primitive operation those tracers pass through gets recorded, and that recording is the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr →, the thing that actually gets compiled and the thing that actually runs on every later call. The same is true under grad and under vmap, not only under jit.

One corollary is easy to miss and explains half the surprises in this track: Python itself is a metaprogramming layer here. The for loops and if statements in your function run once, at trace time, to produce the recording; by the time anything executes on a device, your loops and branches are gone, replaced by whatever primitives they happened to emit.

jit runs your Python once. Everything else follows.

import jax
import jax.numpy as jnp

def attend(q, k):
    return jax.nn.softmax(q @ k.T)

print(jax.make_jaxpr(attend)(jnp.ones((4, 8)), jnp.ones((4, 8))))

Write if x.sum() > 0: inside a jitted function and, if x is traced, JAX raises TracerBoolConversionError, with a message worth reading exactly as it appears: "Attempted boolean conversion of traced array with shape bool[]." The tracer has no value to branch on yet, so there is nothing for if to test.

A print statement inside the same function fires exactly once, at trace time, and what it prints is a tracer, not a number; side effects like that never make it into the recording. Anything your function closes over from the surrounding scope gets baked in as a constant at trace time too, frozen at whatever value it held when tracing happened.

Shapes are part of the trace as well, which is the fact chapter 3 leans on hardest: hand the function an array with a new shape and you get a new trace, not a reused one.

jax.make_jaxpr(f)(*args) shows you the trace directly, without compiling anything, and it is worth running on purpose rather than only reaching for it once something breaks. A jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → has three parts: the invars it takes in, a flat list of equations, each one a primitive applied to earlier values, and the outvars it returns. Every value carries an aval, a shape-and-dtype pair written like f32[4,8].

Look at what tracing an actual softmax produced below and the layer stops being abstract. Two lines of Python, q @ k.T and jax.nn.softmax(...), became twelve equations the machine remembers so you do not have to: the softmax dissolved into seven primitives, the max subtraction carries an explicit stop_gradient that nothing in your source code asked for, the broadcasts Python was hiding are written out in full, and the matmul states its contraction as dimension_numbers instead of leaving it implied.

This chapter only needs fluency at this surface. The kernel path's jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → chapter descends further, into how this same recording lowers to StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → and then to XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla →, but reading equations like the ones below on sight is the skill this path asks of you.

{ lambda ; a:f32[4,8] b:f32[4,8]. let
    c:f32[8,4] = transpose[permutation=(1, 0)] b
    d:f32[4,4] = dot_general[
      dimension_numbers=(([1], [0]), ([], []))
      preferred_element_type=float32
    ] a c
    e:f32[4] = reduce_max[axes=(1,)] d
    f:f32[4] = max -inf e
    g:f32[4,1] = broadcast_in_dim[
      broadcast_dimensions=(0,)
      shape=(4, 1)
      sharding=None
    ] f
    h:f32[4,1] = stop_gradient g
    i:f32[4,4] = sub d h
    j:f32[4,4] = exp i
    k:f32[4] = reduce_sum[axes=(1,)] j
    l:f32[4,1] = broadcast_in_dim[
      broadcast_dimensions=(0,)
      shape=(4, 1)
      sharding=None
    ] k
    m:f32[4,4] = div j l
  in (m,) }

None of this means you are blind to values. jax.debug.print("x = {x}", x=x) stages an actual print into the compiled program, so it runs on every call, with real numbers, unlike a bare Python print that only ever sees a tracer once.

jax.debug.callback generalizes the same idea to any host function, not just printing. And when a NaN shows up somewhere in a long jitted computation and you cannot tell where, jax.config.update("jax_debug_nans", True) re-runs the failing operation eagerly, outside the trace, so the exact op that produced it is the one that raises.

import jax
import jax.numpy as jnp

@jax.jit
def step(x):
    jax.debug.print("x = {x}", x=x)   # runs every call, with real values
    return x * 2

step(jnp.arange(3.0))