What is tracing in JAX?

Tracing is how every JAX transform (jit, grad, vmap, make_jaxpr) inspects your function: it runs your Python code once, but instead of real arrays it passes in tracer objects — stand-ins that carry only the shape and dtype of an array, not its values. As the docs put it, tracers "are stand-ins that encode the shape and dtype of the arrays, but are agnostic to the values." Every JAX operation your function calls on a tracer gets recorded rather than executed, building up a list of primitive operations. Because the tracer has no concrete value, the recorded program is valid for any input of the same shape and dtype, which is exactly what lets JAX compile once and reuse the result.

A jaxpr is JAX’s internal intermediate representation — the small typed program that tracing produces. The docs describe jaxprs as "JAX’s internal intermediate representation (IR) of programs," which are "explicitly typed, functional, first-order, and in algebraic normal form (ANF)." Concretely it is a list of input variables (the invars, one per traced argument), a sequence of equations each applying one primitive (add, mul, sin, dot, reduce_sum…) to previous variables, and a list of output variables. You can see the jaxpr for any function with jax.make_jaxpr(f)(*args). Every transform — compilation, differentiation, batching — is a rule for interpreting or rewriting this one small object, which is why understanding jaxprs is understanding JAX.

Why does my jitted JAX function recompile (retrace)?

Because jit caches compiled code keyed on the shape and dtype of your arguments, and any argument with a new shape or dtype forces a fresh trace and compile. Tracing threw away the values but kept the shape and dtype, so the compiled program is only valid for that exact shape signature; call the function on a differently shaped array and JAX has to trace and compile again. The classic performance bug is a function called in a loop on arrays whose shape changes every iteration (a growing sequence, a variable batch) — it retraces every call and never benefits from compilation. Marking a rarely-changing argument as static (static_argnums / static_argnames) bakes its concrete value into the cache key instead, which trades a recompile-on-change for the ability to use that value in Python-level control flow.

Why can’t I use a Python if statement in a jitted JAX function?

Because under jit the value being tested is a tracer, not a concrete boolean, and Python cannot branch on it. When you write if x < 0 inside a jitted function, x < 0 evaluates to an abstract boolean array that stands for the set {True, False}; when Python tries to coerce that to a real True or False to pick a branch, JAX raises a TracerBoolConversionError — "This concrete value was not available in Python because it depends on the value of the argument." Tracing runs the function once and records a single straight-line path, so a value-dependent branch has no single answer to record. The fixes: use jax.lax.cond (or jnp.where) to record the choice as a data-dependent primitive that both branches trace into, or mark the argument static so its value is concrete at trace time and the ordinary Python if runs at compile time.

What is a pytree in JAX and why does it matter for tracing?

A pytree is any nested structure of containers — dicts, lists, tuples, dataclasses — with arrays at the leaves; a model’s parameters are the canonical example. It matters because jaxprs have no tuple or dict types: JAX flattens every structured input into a flat list of array leaves plus a treedef (the shape of the container), traces over the flat leaves, and unflattens the outputs back into the same structure. The docs note that a function taking a pair produces "an identical jaxpr" to the same function taking two separate arguments, because both flatten to the same two leaves. This flatten-over-leaves / unflatten-with-treedef contract is why you can jit or grad a function whose argument is a whole nested parameter dict without ever unpacking it by hand — and why every transform composes: they all agree to speak in flat leaves.

Tracing → jaxpr: the One Trick Behind Every JAX Transform

Q: What is a pytree in JAX and why does it matter for tracing?

A pytree is any nested structure of containers — dicts, lists, tuples, dataclasses — with arrays at the leaves; a model’s parameters are the canonical example. It matters because jaxprs have no tuple or dict types: JAX flattens every structured input into a flat list of array leaves plus a treedef (the shape of the container), traces over the flat leaves, and unflattens the outputs back into the same structure. The docs note that a function taking a pair produces "an identical jaxpr" to the same function taking two separate arguments, because both flatten to the same two leaves. This flatten-over-leaves / unflatten-with-treedef contract is why you can jit or grad a function whose argument is a whole nested parameter dict without ever unpacking it by hand — and why every transform composes: they all agree to speak in flat leaves.

jit, grad, and vmap look like magic, but they all do the same mundane thing first: they run your Python once, feeding it stand-in values that carry only a shape and a dtype, and record every operation into a small typed program called a jaxpr. Watch a real function become its jaxpr line by line, see the print that fires exactly once, learn why a new input shape makes jit retrace, why a plain Python if throws under jit, and how the pytree flatten/unflatten contract lets all of it compose. Every other JAX page rests on this one.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of Tracing → jaxpr: the One Trick Behind Every JAX Transform — built to be understood, not skimmed.

Questions

What is tracing in JAX?: Tracing is how every JAX transform (jit, grad, vmap, make_jaxpr) inspects your function: it runs your Python code once, but instead of real arrays it passes in tracer objects — stand-ins that carry only the shape and dtype of an array, not its values. As the docs put it, tracers "are stand-ins that encode the shape and dtype of the arrays, but are agnostic to the values." Every JAX operation your function calls on a tracer gets recorded rather than executed, building up a list of primitive operations. Because the tracer has no concrete value, the recorded program is valid for any input of the same shape and dtype, which is exactly what lets JAX compile once and reuse the result.
What is a jaxpr?: A jaxpr is JAX’s internal intermediate representation — the small typed program that tracing produces. The docs describe jaxprs as "JAX’s internal intermediate representation (IR) of programs," which are "explicitly typed, functional, first-order, and in algebraic normal form (ANF)." Concretely it is a list of input variables (the invars, one per traced argument), a sequence of equations each applying one primitive (add, mul, sin, dot, reduce_sum…) to previous variables, and a list of output variables. You can see the jaxpr for any function with jax.make_jaxpr(f)(*args). Every transform — compilation, differentiation, batching — is a rule for interpreting or rewriting this one small object, which is why understanding jaxprs is understanding JAX.
Why does my jitted JAX function recompile (retrace)?: Because jit caches compiled code keyed on the shape and dtype of your arguments, and any argument with a new shape or dtype forces a fresh trace and compile. Tracing threw away the values but kept the shape and dtype, so the compiled program is only valid for that exact shape signature; call the function on a differently shaped array and JAX has to trace and compile again. The classic performance bug is a function called in a loop on arrays whose shape changes every iteration (a growing sequence, a variable batch) — it retraces every call and never benefits from compilation. Marking a rarely-changing argument as static (static_argnums / static_argnames) bakes its concrete value into the cache key instead, which trades a recompile-on-change for the ability to use that value in Python-level control flow.
Why can’t I use a Python if statement in a jitted JAX function?: Because under jit the value being tested is a tracer, not a concrete boolean, and Python cannot branch on it. When you write if x < 0 inside a jitted function, x < 0 evaluates to an abstract boolean array that stands for the set {True, False}; when Python tries to coerce that to a real True or False to pick a branch, JAX raises a TracerBoolConversionError — "This concrete value was not available in Python because it depends on the value of the argument." Tracing runs the function once and records a single straight-line path, so a value-dependent branch has no single answer to record. The fixes: use jax.lax.cond (or jnp.where) to record the choice as a data-dependent primitive that both branches trace into, or mark the argument static so its value is concrete at trace time and the ordinary Python if runs at compile time.
What is a pytree in JAX and why does it matter for tracing?: A pytree is any nested structure of containers — dicts, lists, tuples, dataclasses — with arrays at the leaves; a model’s parameters are the canonical example. It matters because jaxprs have no tuple or dict types: JAX flattens every structured input into a flat list of array leaves plus a treedef (the shape of the container), traces over the flat leaves, and unflattens the outputs back into the same structure. The docs note that a function taking a pair produces "an identical jaxpr" to the same function taking two separate arguments, because both flatten to the same two leaves. This flatten-over-leaves / unflatten-with-treedef contract is why you can jit or grad a function whose argument is a whole nested parameter dict without ever unpacking it by hand — and why every transform composes: they all agree to speak in flat leaves.

Tracing → jaxpr: the One Trick Behind Every JAX Transform

Questions

Related explainers