What actually happens when you call a jitted JAX function the second time?

It goes through a C++ callable rather than through Python. jax.jit hands your function to a C++ PjitFunction along with a Python cache_miss callback; on every call the C++ side parses the arguments into a CallSignature and looks that up in an LRU cache of compiled entries. The signature holds the static arguments, a pytree definition per dynamic argument, the shape and dtype of every array, their shardings, their layouts, whether each is committed to a device, and a snapshot of the 51 configuration states JAX marks as belonging in the jit key. On a hit, the C++ side dispatches straight to the stored XLA executable and Python is never entered. On a miss it calls cache_miss, which runs the Python path and then hands back a MeshExecutableFastpathData object for the C++ side to remember, so the next identical call hits.

Why does my jitted function recompile when nothing changed?

Usually because something the key compares changed and you did not think of it as an input. Four causes account for most of it. A wrapper built inline, as in jit(lambda x: f(x, y)), is a different function object every call, and the caches key on that object, so nothing ever hits. Entering a mesh context manager changes the context mesh, which is part of the params cache key and part of the C++ configuration snapshot. Flipping a config flag that is marked as belonging in the jit key does the same, and there are 51 of those. And an argument whose shape or dtype moved, even by one element, changes the abstract values. JAX will tell you which one: set JAX_EXPLAIN_CACHE_MISSES=1 and the tracing cache prints TRACING CACHE MISS at your call site, followed by the smallest difference between this key and the closest previous one, ranked so a context difference is reported before a pytree difference.

What is the difference between the in-memory jit cache and the persistent compilation cache?

They are keyed on completely different things. The in-memory caches are keyed on Python objects: the function, the argument signature, the abstract values, the jaxpr, the lowered module object. They are per-process and they die with it. The persistent cache is content-addressed: its key is a sha256 over the StableHLO bytecode, the jaxlib version, the backend platform and version, the XLA flags that affect compilation, the serialized compile options, and the PJRT topology fingerprint. Because it hashes the module rather than the Python identity of anything, two processes, or one process before and after a restart, share entries as long as the lowered module comes out byte-identical. It is off unless you set jax_compilation_cache_dir.

Does a new static argument value force a full recompile?

Every in-memory layer misses, yes. Static arguments live in the ArgumentSignature the C++ side builds, compared by type identity and then by equality, so a new value is a new signature: the C++ cache misses, the params cache misses, tracing runs, lowering runs. The jit docstring is explicit that calling with different values for these constants triggers recompilation. What survives is the on-disk layer, and only sometimes. If the new static value changes the traced program, the module differs and the persistent key differs too. If it does not, for example a static flag your function never branches on, the module comes out identical and the persistent cache answers with the executable it already has, so you pay for the retrace and the lowering but not for the XLA compile.

How does JAX run an operation with no jit around it at all?

By compiling it, through the same machinery. There is no eager kernel registry. A primitive dispatched outside any jit goes to apply_primitive, which asks xla_primitive_callable for a callable and calls it; that function builds a tiny wrapper around the primitive bind and returns api.jit of it. It is decorated with util.cache, so each primitive plus parameter combination compiles once per process and every later call rides the same jit fast path as your own functions. apply_primitive also swaps jax_disable_jit off around the call, so the compile happens even for a user who disabled jit globally.

What do the three JAX_LOG_COMPILES timing lines mean?

Each one marks a different expensive station. Finished tracing NAME for jit is the trace: running your Python once on abstract values to build a jaxpr. Finished jaxpr to MLIR module conversion NAME is the lowering: turning that jaxpr into a StableHLO module. Finished XLA compilation of NAME is the compile: handing the module to the backend and getting an executable. Setting JAX_LOG_COMPILES=1 raises these three from debug level to warning so they show up without configuring logging. Reading which of the three appeared tells you how far down the stack the call fell before something answered, which is a faster diagnosis than guessing.

Dispatch and Compilation: The Six Caches Under a jit Call

The first call to a jitted function takes a second. The second takes microseconds. Nothing about your code changed, so the difference is entirely bookkeeping: six caches stacked one under the other, each one catching a different kind of repetition, each with its own key. This page walks the stack from the C++ fast path down to the executable on disk, naming every cache by the variable that holds it and every key by the fields it actually compares, at one pinned commit of jax-ml/jax. Then it puts the stack under your hands: issue calls, change one thing at a time, and watch which cache answers and which ones never get asked. By the end, a surprise recompile is a question with a procedure, not a mystery.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of Dispatch and Compilation: The Six Caches Under a jit Call — built to be understood, not skimmed.

Questions

What actually happens when you call a jitted JAX function the second time?: It goes through a C++ callable rather than through Python. jax.jit hands your function to a C++ PjitFunction along with a Python cache_miss callback; on every call the C++ side parses the arguments into a CallSignature and looks that up in an LRU cache of compiled entries. The signature holds the static arguments, a pytree definition per dynamic argument, the shape and dtype of every array, their shardings, their layouts, whether each is committed to a device, and a snapshot of the 51 configuration states JAX marks as belonging in the jit key. On a hit, the C++ side dispatches straight to the stored XLA executable and Python is never entered. On a miss it calls cache_miss, which runs the Python path and then hands back a MeshExecutableFastpathData object for the C++ side to remember, so the next identical call hits.
Why does my jitted function recompile when nothing changed?: Usually because something the key compares changed and you did not think of it as an input. Four causes account for most of it. A wrapper built inline, as in jit(lambda x: f(x, y)), is a different function object every call, and the caches key on that object, so nothing ever hits. Entering a mesh context manager changes the context mesh, which is part of the params cache key and part of the C++ configuration snapshot. Flipping a config flag that is marked as belonging in the jit key does the same, and there are 51 of those. And an argument whose shape or dtype moved, even by one element, changes the abstract values. JAX will tell you which one: set JAX_EXPLAIN_CACHE_MISSES=1 and the tracing cache prints TRACING CACHE MISS at your call site, followed by the smallest difference between this key and the closest previous one, ranked so a context difference is reported before a pytree difference.
What is the difference between the in-memory jit cache and the persistent compilation cache?: They are keyed on completely different things. The in-memory caches are keyed on Python objects: the function, the argument signature, the abstract values, the jaxpr, the lowered module object. They are per-process and they die with it. The persistent cache is content-addressed: its key is a sha256 over the StableHLO bytecode, the jaxlib version, the backend platform and version, the XLA flags that affect compilation, the serialized compile options, and the PJRT topology fingerprint. Because it hashes the module rather than the Python identity of anything, two processes, or one process before and after a restart, share entries as long as the lowered module comes out byte-identical. It is off unless you set jax_compilation_cache_dir.
Does a new static argument value force a full recompile?: Every in-memory layer misses, yes. Static arguments live in the ArgumentSignature the C++ side builds, compared by type identity and then by equality, so a new value is a new signature: the C++ cache misses, the params cache misses, tracing runs, lowering runs. The jit docstring is explicit that calling with different values for these constants triggers recompilation. What survives is the on-disk layer, and only sometimes. If the new static value changes the traced program, the module differs and the persistent key differs too. If it does not, for example a static flag your function never branches on, the module comes out identical and the persistent cache answers with the executable it already has, so you pay for the retrace and the lowering but not for the XLA compile.
How does JAX run an operation with no jit around it at all?: By compiling it, through the same machinery. There is no eager kernel registry. A primitive dispatched outside any jit goes to apply_primitive, which asks xla_primitive_callable for a callable and calls it; that function builds a tiny wrapper around the primitive bind and returns api.jit of it. It is decorated with util.cache, so each primitive plus parameter combination compiles once per process and every later call rides the same jit fast path as your own functions. apply_primitive also swaps jax_disable_jit off around the call, so the compile happens even for a user who disabled jit globally.
What do the three JAX_LOG_COMPILES timing lines mean?: Each one marks a different expensive station. Finished tracing NAME for jit is the trace: running your Python once on abstract values to build a jaxpr. Finished jaxpr to MLIR module conversion NAME is the lowering: turning that jaxpr into a StableHLO module. Finished XLA compilation of NAME is the compile: handing the module to the backend and getting an executable. Setting JAX_LOG_COMPILES=1 raises these three from debug level to warning so they show up without configuring logging. Reading which of the three appeared tells you how far down the stack the call fell before something answered, which is a faster diagnosis than guessing.

Dispatch and Compilation: The Six Caches Under a jit Call

Questions

Related explainers