What tracing takes, exactly

Trace happens once, but what tracing captures from outside the function often catches people off guard. Close over a Python-level array, a weight matrix built outside the traced function, say, and JAX doesn't keep a live reference to it. It freezes the array's value into the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → as a constant, carried alongside the traced computation as what the internals call a constvar. Call the function again with the same shapes and that constant is still the one baked in from the first trace, not read fresh from Python each time.

Python side effects run once at trace time, not once per call. A print statement inside a traced function fires while JAX is walking the function with symbolic tracers standing in for real arrays, and that walk happens exactly once per distinct shape and dtype signature. Every call after that runs the compiled, traced computation directly, and the Python body never executes again. If you're watching stdout to confirm a kernel ran on some later call, there's nothing to see: the print already happened, once, during the original trace.

Shapes and dtypes freeze into the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → at trace time as well. Every traced value carries its exact shape and dtype into the recorded computation, so a new input shape doesn't reuse the trace you already have. It triggers a fresh trace and a fresh compile for that new signature, every time the signature changes: a function traced for one batch size retraces completely for another.

Data-dependent Python branching disappears into the traced path entirely. An ordinary if statement checks a traced value's Python truthiness, and tracing has to resolve that check to a concrete bool right then, so it picks one branch and the other branch never gets recorded at all. That's harmless when the condition is static and known ahead of time, the same value on every call. It's wrong when the condition depends on runtime data the trace can't reduce to a fixed answer: whichever branch got traced is the only one that will ever run, no matter what the actual value turns out to be later.

lax.cond exists for exactly that case. It lowers to a real conditional that the compiled program evaluates at runtime instead of at trace time, so both branches get recorded and the choice between them happens on every call, not once during tracing. The running example in this chapter never needed that, its branching was all static, but plenty of kernels do. Chapter 03 walks through what lax.cond looks like once it's lowered.