Learn AI Create the Future — with Pranay
Main site

JAX and Flax

jit and tracing

jax.jit runs your function once with stand-in values, records every operation, and compiles the recording into fast machine code it replays from then on.

Read these first

On this page 5
  1. Why it exists
  2. How it works — and the catch
  3. A real example you have seen
  4. Remember this
  5. What to learn next

One lesson, three depths. Pick the one that fits you today — you can switch any time.

Beginner — No maths. Plain English.

jax.jit watches your function run once, writes down every maths step it performs, and turns that recording into fast compiled code that replaces your Python.

Think of a stenographer sitting beside you while you fill a long tax form for the first time. You work slowly, deciding each line; the stenographer records the finished procedure. Next year, nobody thinks at all — a clerk executes the recorded procedure at top speed, whatever numbers you bring.

The recording pass is called tracing. The recorded procedure is compiled — translated into fast machine code — by a compiler called XLA. "jit" stands for just in time: the compilation happens at first call, not before.

Why it exists

Python decides everything at the last moment, which makes it friendly and slow. Numerical code suffers most: a training step may run thousands of small operations, each paying Python's decision-making cost.

Tracing removes Python from the loop. Once the procedure is recorded, the compiler sees the whole calculation at once. So it can merge steps, skip waste, and tailor the code to your exact array sizes. Speed-ups of 10x and beyond on small-operation-heavy code are normal.

How it works — and the catch

first call:    Python runs once ──▶ every operation recorded ──▶ compiled
later calls:   compiled code runs, Python skipped

the catch:     during recording, your function sees STAND-INS, not numbers.
               "is x > 0?"  ──▶  the stand-in has no answer  ──▶  error

The stand-in (called a tracer) knows the shape and type of the value — "a decimal number", "a 3-by-3 grid" — but not the value itself. So ordinary maths records fine, but a Python if asking about the value cannot proceed. This one catch explains most JAX error messages you will ever see.

A real example you have seen

A bank's loan-approval checklist. It was designed once, by slow careful committee — every rule written down. Clerks then process thousands of applications fast by following the sheet. But the sheet must cover every case in advance: a clerk cannot phone the committee mid-application. Compiled functions are the sheet; Python is the committee.

Remember this

  • Tracing = one recorded run with stand-in values; then compiled code replays.
  • Stand-ins know shape and type, never the actual numbers.
  • Python if/while on actual values cannot be traced — JAX has its own tools for that.

What to learn next

Developer — Code and libraries.

Setup

bash
pip install jax

Outputs verified with jax 0.6.2, CPU. Timing numbers are environment-dependent — expect different milliseconds, similar ratio.

Measuring the speedup

jit_speed.py
import time
import jax
import jax.numpy as jnp


def step(m):
    for _ in range(50):
        m = (m @ m) / jnp.linalg.norm(m)
    return m


fast = jax.jit(step)
x = jax.random.normal(jax.random.key(0), (200, 200))

fast(x).block_until_ready()                     # first call: tracing + compiling

t0 = time.perf_counter(); step(x).block_until_ready()
plain = time.perf_counter() - t0
t0 = time.perf_counter(); fast(x).block_until_ready()
jitted = time.perf_counter() - t0
print(f"plain: {plain*1000:.1f} ms, jitted: {jitted*1000:.1f} ms")
Output
plain: 65.5 ms, jitted: 18.5 ms

Three details carry the lesson. The warm-up call pays the trace-and-compile cost before any timing — time the first call and you would measure compilation, not execution. .block_until_ready() forces the asynchronous computation to finish; without it you time how fast JAX queues work, which is nearly instant and meaningless. And the Python for loop was fine to trace: its trip count depends on no array value, so it unrolls into 50 recorded steps.

Seeing the tracer

see_tracer.py
import jax
import jax.numpy as jnp


@jax.jit
def f(x):
    print("tracing, x is:", x)
    return x * 2


print(f(jnp.array(1.0)))
print(f(jnp.array(5.0)))     # same shape and dtype: no new trace
Output
tracing, x is: Traced<~float32[]>with<DynamicJaxprTrace>
2.0
10.0

The print fired once, and what it saw was no number — Traced<~float32[]> is the stand-in, carrying dtype and shape (scalar) only. The second call, same shape and dtype, replayed silently. New shape or dtype would trace again: compiled code is specialised per signature, the same cache-by-signature behaviour as tf.function.

Where the catch bites

value_branch.py
import jax


@jax.jit
def relu_wrong(x):
    if x > 0:            # Python asks: is the tracer true or false?
        return x
    return 0.0


try:
    relu_wrong(3.0)
except jax.errors.TracerBoolConversionError as err:
    print("error:", str(err).splitlines()[0][:80])
Output
error: Attempted boolean conversion of traced array with shape bool[].

The tracer cannot answer x > 0 — it has no value. The fixes, in order of preference: jnp.where(x > 0, x, 0.0) computes both sides and selects (perfect here); lax.cond and friends for genuine branching; static_argnums to declare an argument a compile-time constant, retracing per new value.

Common mistakes

Benchmarking the first call. Compilation takes longer than execution — sometimes by seconds. Warm up, then time.

Forgetting block_until_ready() in timings. Async dispatch makes everything look instant. Fence before reading the clock.

Jitting a function called once. Compilation is an investment repaid over many calls. A one-shot function is faster un-jitted.

Retracing on every call. Feeding batches whose length changes each step (ragged text batches, final partial batch) compiles one program per shape. Symptom: training "hangs" at random steps. Fix: pad to fixed shapes — the same discipline as padding sequences.

Side effects that vanish. print, appending to lists, mutating globals — all happen at trace time only. For debugging values inside jitted code, jax.debug.print("x = {}", x) runs at execution time.

Try it yourself

Add jax.debug.print("running, x = {}", x) inside f above, next to the plain print. Call it three times with same-shaped inputs. Predict which line appears how many times, then check.

What to learn next

Researcher — Mathematics and papers.

What tracing constructs

Tracing evaluates the Python function with abstract values (ShapedArray: dtype, shape, no data) and records primitive applications into a jaxpr — a typed, first-order, ANF-style intermediate representation. Inspect it with jax.make_jaxpr(f)(x). The jaxpr is closed over the trace-time-resolved control flow: Python conditionals on static values are gone, loops unrolled, leaving pure dataflow. XLA then lowers it: fusion of elementwise chains into single kernels, buffer reuse from liveness analysis (the machinery that makes immutability cheap), layout selection, and code generation per backend.

Compilation caching is keyed on the abstraction: $\text{key} = (\text{function identity}, \text{tree structure of args}, \text{ShapedArray of each leaf}, \text{static arg values})$. Every distinct key pays trace + compile once — $O(\text{ops})$ trace, compile time in practice seconds-ish for large models — and hits are dictionary lookups. The retracing pathologies of the developer block are exactly cache-key churn.

The specialisation trade

Shape-specialised compilation is the deliberate bet distinguishing JAX from shape-polymorphic eager frameworks. Fixing shapes at compile time lets XLA allocate statically, unroll, vectorise, and pre-plan communication — the performance foundation that made JAX standard in large-scale training (used by DeepMind and Google Research for LLM-scale work). The cost is the padding discipline for dynamic shapes and the two-language feel: Python is the metaprogramming layer; the jaxpr is the program.

jit composes with the other transformations — jit(grad(f)), vmap(jit(f)) — because all of them are jaxpr-to-jaxpr or trace-level rewrites; composability of transformations is the core design claim of Frostig, Johnson and Leary (2018), Compiling machine learning programs via high-level tracing.

Donation (donate_argnums) extends buffer liveness analysis across the call boundary, allowing input buffers to be reused for outputs — the standard trick for updating large model states in place under a functional API.

References

  • Frostig, Johnson, Leary (2018), Compiling machine learning programs via high-level tracing.
  • Sabne (2020), XLA: compiling machine learning for peak performance.
  • Bradbury et al. (2018–), JAX: composable transformations of Python+NumPy programs — the library itself and its documentation on tracing semantics.

What to learn next