Jit

Without jit, JAX dispatches one primitive at a time: each op runs, materializes its result, and hands that result to the next op. With jit, XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → sees the whole jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → at once before any of it runs, so it can fuse chains of elementwise operations together, keep intermediates in registers or cache instead of round-tripping through memory, and emit a single executable for the whole function.

The win from fusing compounds as the program grows, since a longer jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → gives XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → more to fuse. The cost sits on the other side of the ledger, paid once per shape signature rather than once per call: only the first call for a given signature pays trace time plus compile time.

Every jitted function is backed by a cache, keyed on four things at once: function identity, the pytree structure of the arguments, the shape and dtype of every leaf, and the value of any argument marked static. Match all four and you get the cached executable back in microseconds. Change any one of them, a new shape, a new dtype, a different pytree structure, and JAX retraces and recompiles from scratch; a new static value forces exactly the same retrace, because a static argument is not really an input to the compiled program, it is baked into the trace like a closed-over constant.

Three habits account for almost every recompile bug people report: passing a fresh lambda or closure on every call, which counts as a new function identity even though the code inside is identical; marking a high-cardinality argument static, so nearly every call earns its own executable; and Python scalars flipping between weak and strong dtypes across calls, which changes the key without a line of code that looks like it should. static_argnums and static_argnames are the tool for moving an argument out of the traced set on purpose, when you want its value baked in rather than treated as data.

A recompile is never mysterious. You changed the key.

from functools import partial

import jax
import jax.numpy as jnp

@partial(jax.jit, static_argnames="k")
def topk_sum(x, k):
    return jax.lax.top_k(x, k)[0].sum()

x = jnp.arange(100.0)
topk_sum(x, 3)   # trace + compile, k = 3 baked in
topk_sum(x, 3)   # cache hit
topk_sum(x, 5)   # new static value: trace + compile again

Some computations cannot be traced at all, not because JAX is missing a feature but because the question itself has no answer at trace time. x[x > 0] is the standard example: the output shape depends on how many entries of x are positive, and that count does not exist until the values do. Two escapes exist: pick a static upper bound and mask down to it, or pull that one step outside the jitted region entirely and let it run as ordinary, untraced Python.

Data-dependent Python control flow runs into the same wall from a different angle. A branch whose condition depends on a traced value needs restructuring, either with the lax control-flow primitives chapter 6 covers, or by making the branching value itself static so the choice resolves at trace time instead of run time.

Compilation does not have to happen at call time. jax.jit(f).lower(*args) produces the StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → for a function before anything runs at all; .compile() turns that into an executable; and compiled.cost_analysis() returns the compiler's own estimate of flops and bytes moved, no execution required. Reading the cost of a computation before running it is a habit worth owning on its own, not only a debugging trick.

The numbers below are real, not illustrative: compiled.cost_analysis()["flops"] on a 256x256 matmul reports 33554432.0, which is exactly 2 * 256**3, the matmul's true FLOP count with nothing hidden. On JAX versions before 0.5, cost_analysis() returned a one-element list wrapping this same dict; the dict form shown here is current as of 0.4.38, and if your installed version differs, that is the first thing to check.

jax.export takes the same lowered form one step further, serializing it so a function can be compiled once and served elsewhere, outside the Python process that traced it.

import jax
import jax.numpy as jnp

lowered = jax.jit(lambda x: x @ x).lower(jnp.ones((256, 256)))
print(lowered.as_text()[:300])   # StableHLO, before anything runs
compiled = lowered.compile()
print(compiled.cost_analysis()["flops"])   # the compiler's own estimate: 33554432.0