What goes into the JAX compilation cache key?

Eight named entries, hashed in a fixed order into one SHA-256, and the order is frozen because a running hash is order-sensitive. First the computation, meaning the MLIR module serialized to bytecode after canonicalisation. Second the jaxlib version string, which is why every jaxlib upgrade invalidates the whole persistent cache. Third the backend, hashed as two strings, its platform name then its platform version. Fourth the XLA flags gathered from XLA_FLAGS, LIBTPU_INIT_ARGS and the command line, sorted and filtered. Fifth the compile options proto, deep-copied and with its debug fields flattened first. Sixth the accelerator configuration, which is a topology fingerprint from the runtime hashed as exactly eight big-endian bytes, falling back to one device_kind string per device when the backend cannot produce a topology. Seventh the compression algorithm name, in practice zstandard when the zstandard module is importable and zlib otherwise. Eighth a custom_hook that ships returning the empty string and therefore contributes zero bytes unless somebody replaces it. The returned key is the module name, a dash, and 64 hex characters; the name is a prefix for human readability and is not itself hashed, because it is already inside the bytecode hashed as entry one.

Does setting a dump flag change the JAX cache key?

No, and that is deliberate. Fourteen of the twenty-six entries on the exclusion list are dump flags: write the HLO as text, as a proto, as a dot graph, as HTML, into a directory, with a timestamp, and so on. When the flag arrives as a string, each collected flag is split on its first equals sign and its name tested against that list; a match is logged as not included and skipped. When the same thing arrives as a typed field on the compile options object, it is not covered by the list at all, so the file zeroes twenty such fields by hand on a deep copy before serializing the proto. Both places are wrapped in matching LINT.IfChange markers so neither is edited alone. The practical consequence is that turning on a dump to investigate a slow compile does not throw away the cache entry you are investigating. The reverse would be much worse: if a dump flag did move the key, every engineer debugging a model would silently be compiling a second copy of it.

Why do two identical JAX functions get different cache keys?

Usually because metadata is in the key. By default jax_compilation_cache_include_metadata_in_key is false, so the module is canonicalised before hashing: it is cloned, run through the MLIR strip-debuginfo pass, and only then serialized. File names and line numbers are gone, so the same function defined at a different source line, or moved to another file, still hits. Turn the flag on and those two definitions get different keys, which the test suite asserts in both directions. The cost of leaving it off is stated in the flag help: an executable loaded from the cache may carry stale metadata that shows up in profiles. There is a second, narrower version of the same problem on TPU. A Pallas kernel sits in the module as a whole serialized MLIR module, base64-encoded inside a JSON blob inside one attribute, and an ordinary MLIR pass cannot see into a string, so a separate helper decodes it, strips debug info from the kernel itself, minifies the JSON and writes it back. Without that step, two checkouts in different directories would never share a cache entry for the same kernel.

Why does a multi-process GPU job compile the same program once per process?

It does not, as long as this path is working. Each process in a multi-process job holds a different device assignment, and hashing it would give every process a different key for one compilation. So the cache key builder passes strip_device_assignment when the backend platform is gpu, and the compile-options hasher then reads the replica count and computation count off the assignment and rebuilds it from np.arange reshaped to those two counts. The real device numbers are replaced by 0, 1, 2 and so on, so every process serializes the same bytes and lands on the same key. Note the condition is the platform string, not the process count: on a single-process GPU run the assignment is flattened too, and on TPU it is not flattened at all, which is why the test for this asserts equal keys on gpu and unequal keys everywhere else.

How do I debug a JAX compilation cache miss?

Turn on debug logging for the jax._src.cache_key module and diff two runs. The key builder calls a logging companion after every one of the eight entries, and when debug logging is enabled that companion prints two lines per entry: the digest of that entry hashed alone into a fresh SHA-256, and the running digest of the real hash object as it stands after that entry. So one key produces sixteen lines in a fixed order with the entry names attached. Diff the traces from a run that hit and a run that missed, and the first line that disagrees names the ingredient that moved. Two caveats. The per-entry digest is produced by running the same hashing lambda a second time, so under debug logging the module is canonicalised and serialized twice, which is not free on a large program; that is why the whole body sits behind an isEnabledFor check. And the running digest is cumulative, so once one entry moves, every line below it moves too. The first disagreement is the answer; the ones under it are consequences.

What JAX Hashes Before It Decides Not to Compile: cache_key.py, Line by Line

Every jitted call ends at one question: have I compiled this exact thing before? Answering it means turning a compiler module, a set of devices, a pile of options and the environment your process happens to be in into one string. This is the 399-line file that does it, read whole at a pinned commit. Eight named ingredients go into a single running SHA-256, in a frozen order, and roughly a third of the file exists only to make the bytes arriving at those updates identical on two machines doing the same work: clone before mutating, strip debug info, then strip the debug info hiding inside a base64 string that the strip pass walks straight past, replace an unstable pointer with a constant, renumber the devices so a multi-process GPU job shares one key. Then twenty-six flag names are deliberately dropped, because they change what lands on your disk rather than what lands in the binary. The signature exhibit is the entry list itself, ported in order, with real digests you perturb one ingredient at a time; a second exhibit ports the flag filter so you can watch a dump flag get sorted in and then thrown away.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of What JAX Hashes Before It Decides Not to Compile: cache_key.py, Line by Line — built to be understood, not skimmed.

Questions

What goes into the JAX compilation cache key?: Eight named entries, hashed in a fixed order into one SHA-256, and the order is frozen because a running hash is order-sensitive. First the computation, meaning the MLIR module serialized to bytecode after canonicalisation. Second the jaxlib version string, which is why every jaxlib upgrade invalidates the whole persistent cache. Third the backend, hashed as two strings, its platform name then its platform version. Fourth the XLA flags gathered from XLA_FLAGS, LIBTPU_INIT_ARGS and the command line, sorted and filtered. Fifth the compile options proto, deep-copied and with its debug fields flattened first. Sixth the accelerator configuration, which is a topology fingerprint from the runtime hashed as exactly eight big-endian bytes, falling back to one device_kind string per device when the backend cannot produce a topology. Seventh the compression algorithm name, in practice zstandard when the zstandard module is importable and zlib otherwise. Eighth a custom_hook that ships returning the empty string and therefore contributes zero bytes unless somebody replaces it. The returned key is the module name, a dash, and 64 hex characters; the name is a prefix for human readability and is not itself hashed, because it is already inside the bytecode hashed as entry one.
Does setting a dump flag change the JAX cache key?: No, and that is deliberate. Fourteen of the twenty-six entries on the exclusion list are dump flags: write the HLO as text, as a proto, as a dot graph, as HTML, into a directory, with a timestamp, and so on. When the flag arrives as a string, each collected flag is split on its first equals sign and its name tested against that list; a match is logged as not included and skipped. When the same thing arrives as a typed field on the compile options object, it is not covered by the list at all, so the file zeroes twenty such fields by hand on a deep copy before serializing the proto. Both places are wrapped in matching LINT.IfChange markers so neither is edited alone. The practical consequence is that turning on a dump to investigate a slow compile does not throw away the cache entry you are investigating. The reverse would be much worse: if a dump flag did move the key, every engineer debugging a model would silently be compiling a second copy of it.
Why do two identical JAX functions get different cache keys?: Usually because metadata is in the key. By default jax_compilation_cache_include_metadata_in_key is false, so the module is canonicalised before hashing: it is cloned, run through the MLIR strip-debuginfo pass, and only then serialized. File names and line numbers are gone, so the same function defined at a different source line, or moved to another file, still hits. Turn the flag on and those two definitions get different keys, which the test suite asserts in both directions. The cost of leaving it off is stated in the flag help: an executable loaded from the cache may carry stale metadata that shows up in profiles. There is a second, narrower version of the same problem on TPU. A Pallas kernel sits in the module as a whole serialized MLIR module, base64-encoded inside a JSON blob inside one attribute, and an ordinary MLIR pass cannot see into a string, so a separate helper decodes it, strips debug info from the kernel itself, minifies the JSON and writes it back. Without that step, two checkouts in different directories would never share a cache entry for the same kernel.
Why does a multi-process GPU job compile the same program once per process?: It does not, as long as this path is working. Each process in a multi-process job holds a different device assignment, and hashing it would give every process a different key for one compilation. So the cache key builder passes strip_device_assignment when the backend platform is gpu, and the compile-options hasher then reads the replica count and computation count off the assignment and rebuilds it from np.arange reshaped to those two counts. The real device numbers are replaced by 0, 1, 2 and so on, so every process serializes the same bytes and lands on the same key. Note the condition is the platform string, not the process count: on a single-process GPU run the assignment is flattened too, and on TPU it is not flattened at all, which is why the test for this asserts equal keys on gpu and unequal keys everywhere else.
How do I debug a JAX compilation cache miss?: Turn on debug logging for the jax._src.cache_key module and diff two runs. The key builder calls a logging companion after every one of the eight entries, and when debug logging is enabled that companion prints two lines per entry: the digest of that entry hashed alone into a fresh SHA-256, and the running digest of the real hash object as it stands after that entry. So one key produces sixteen lines in a fixed order with the entry names attached. Diff the traces from a run that hit and a run that missed, and the first line that disagrees names the ingredient that moved. Two caveats. The per-entry digest is produced by running the same hashing lambda a second time, so under debug logging the module is canonicalised and serialized twice, which is not free on a large program; that is why the whole body sits behind an isEnabledFor check. And the running digest is cumulative, so once one entry moves, every line below it moves too. The first disagreement is the answer; the ones under it are consequences.

What JAX Hashes Before It Decides Not to Compile: cache_key.py, Line by Line

Questions

Related explainers