How do I print the jaxpr of a jitted function in JAX?

Two spellings give you the same object. jax.make_jaxpr(f)(x) is the short one and jax.jit(f).trace(x).jaxpr is the long one, and they are the same code path: the body of make_jaxpr builds a jitted version of your function, calls trace on it, and returns the jaxpr that comes out, reattaching any hoisted constants. What you get back is the set of equations your function performed, with the values removed. A jaxpr for sin of cos of a scalar reads as one input binder, two equations each binding a new typed name to the result of one primitive, and one output tuple. Pass return_shape=True and you get a pair instead, the jaxpr plus a pytree of the output shapes and dtypes. Stopping at the traced object rather than at make_jaxpr buys three extra properties: the output info without lowering anything, which is exactly what jax.eval_shape returns; the set of effects that survived the trace; and per-argument shape, dtype and donation information. Neither spelling needs an accelerator, because nothing has been asked of a device yet.

What is the difference between jax.jit(f).lower() and .compile()?

Lowering turns the jaxpr into one MLIR module in the StableHLO dialect, with sharding decisions attached as attributes. Compiling hands that module to XLA and gets an executable back. The artefacts differ accordingly. On the lowered object, as_text() prints StableHLO by default and the same module in HLO text form if you pass the string hlo; any other dialect string raises a ValueError naming what you passed. A debug_info keyword decides whether source locations ride along, which in practice is what makes a large module readable. Its cost analysis carries an explicit warning in its own docstring that it estimates execution cost in the absence of compiler optimizations, which may drastically affect the cost. On the compiled object, as_text() reaches through to the runtime executable for the optimised HLO, cost_analysis() is the post-optimisation twin, memory_analysis() asks the executable for compiled memory statistics, and the sharding and format properties say where the compiler decided each input has to live. Every one of the compiled artefacts catches a not-implemented error and returns None rather than raising, so an empty result means a backend that cannot answer, not a small number.

Why does print() inside a jitted JAX function only run once?

Because your Python body only runs once. Every jitted call assembles a key from the argument signature and the abstract input types and looks it up in a cache; on a hit the cached parameters come straight back and the body is never touched, and only a miss traces. So the print fires at trace time, which happens once per distinct key. The key has four moving parts worth knowing. It holds a weak reference to the function object itself, so re-running the cell that defines your function produces a new object and a fresh miss. Static arguments are part of the argument signature. The abstract types carry shape, dtype and sharding, so a new shape misses. And a session-wide tracing context rides on top, so flipping a config anywhere invalidates every entry. To make the print fire again, clear the cache for that function, which evicts the weak reference and then clears the shared parameter cache wholesale. To print on every call instead, use jax.debug.print, which stages a callback into the compiled program; its docstring warns that it does not work with f-strings because formatting is delayed, so the format string and the values are passed separately.

How do I dump the StableHLO or HLO that JAX produces?

Two options, and both must be set before JAX is imported, because both are backed by a holder that reads the environment once at import and has no context manager. JAX_DUMP_IR_TO names a directory, and an empty value means no dumping at all. JAX_DUMP_IR_MODES is a comma-separated list defaulting to stablehlo, and the other three values are jaxpr, jaxpr_html and eqn_count_pprof. Filenames are predictable, which is what makes two runs diffable: modules come out as a four-digit counter, the module name with unsafe characters stripped, and a stage name; jaxprs as a six-digit counter and the function name, with an extension naming the form. The directory is created for you, and the special value sponge redirects everything to a test output directory named by the environment. One detail catches people out: the only site that dumps a module with a stage name sits at the top of the compile entry point, before the persistent cache is consulted, so the file appears even when the compile was served from cache. The before-and-after-pass dumping that XLA does is a separate system driven by XLA flags, and turning it on does not change the compilation cache key, because the dump flags sit on an explicit exclusion list.

Why is my JAX compilation cache directory empty?

Almost always because your compiles are too fast to be worth caching. Three gates sit between a successful compile and a file on disk. The write path times the compile and compares it against a minimum, whose default is one second, and skips anything faster on the reasoning that recompiling is cheaper than reading. A second gate is a minimum entry size in bytes, defaulting to zero but documented as something the runtime may raise to suit the filesystem. A third restricts writes to process zero to avoid contention on shared filesystems, which is free in a notebook but explains sparse caches in multi-host jobs. Every skip logs a line explaining itself, and by default those lines sit at debug level where nobody sees them; setting the cache-miss explanation option raises them to warning level, at which point the empty directory explains itself. Two more things worth separating: the enable switch already defaults to true, so a missing directory setting rather than a disabled cache is the usual configuration error, and clearing the JAX caches does not touch the persistent one, as the docstring of that function says outright.

Writing JAX Notebooks That Show Their Work

A cell that prints a correct array proves nothing about what JAX did to produce it. This reads the machinery at a pinned commit of jax-ml/jax and lays out what a notebook can put on the record. A jitted call is four stations, not one step: the wrapper, the traced form, the lowered module, the compiled executable, each a plain object you can hold in a variable and interrogate, and the page walks them as a scene whose artefact list at every stop is read from the real method sets. What make_jaxpr actually is, which is the trace station under a friendlier name. The two dialects a lowering will print, and the ValueError for a third. Why the cost analysis before compiling says outright that it is looking at unoptimised code, and which one to read instead. The three stage timers JAX already publishes on a public event bus, with their verbatim log strings, so a nine-second cell becomes a trace time, a lowering time and a compile time. Then the case worth committing to: the print inside your function that fires once and goes quiet, walked call by call against the real cache-key ingredients, including the redefinition a re-run cell causes. The 125 configuration options split into the 115 you can scope and the ten you get one shot at, the three that raise if a backend is already up, where the IR dumps land and what the filenames mean, why a cache directory you configured correctly stays empty, and the five ordered categories JAX will use to explain its own cache miss, including the message it keeps for a function being redefined on the same line. Closing with the eight-cell order that survives a runtime restart, and what breaks when each cell is missing.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of Writing JAX Notebooks That Show Their Work — built to be understood, not skimmed.

Questions

How do I print the jaxpr of a jitted function in JAX?: Two spellings give you the same object. jax.make_jaxpr(f)(x) is the short one and jax.jit(f).trace(x).jaxpr is the long one, and they are the same code path: the body of make_jaxpr builds a jitted version of your function, calls trace on it, and returns the jaxpr that comes out, reattaching any hoisted constants. What you get back is the set of equations your function performed, with the values removed. A jaxpr for sin of cos of a scalar reads as one input binder, two equations each binding a new typed name to the result of one primitive, and one output tuple. Pass return_shape=True and you get a pair instead, the jaxpr plus a pytree of the output shapes and dtypes. Stopping at the traced object rather than at make_jaxpr buys three extra properties: the output info without lowering anything, which is exactly what jax.eval_shape returns; the set of effects that survived the trace; and per-argument shape, dtype and donation information. Neither spelling needs an accelerator, because nothing has been asked of a device yet.
What is the difference between jax.jit(f).lower() and .compile()?: Lowering turns the jaxpr into one MLIR module in the StableHLO dialect, with sharding decisions attached as attributes. Compiling hands that module to XLA and gets an executable back. The artefacts differ accordingly. On the lowered object, as_text() prints StableHLO by default and the same module in HLO text form if you pass the string hlo; any other dialect string raises a ValueError naming what you passed. A debug_info keyword decides whether source locations ride along, which in practice is what makes a large module readable. Its cost analysis carries an explicit warning in its own docstring that it estimates execution cost in the absence of compiler optimizations, which may drastically affect the cost. On the compiled object, as_text() reaches through to the runtime executable for the optimised HLO, cost_analysis() is the post-optimisation twin, memory_analysis() asks the executable for compiled memory statistics, and the sharding and format properties say where the compiler decided each input has to live. Every one of the compiled artefacts catches a not-implemented error and returns None rather than raising, so an empty result means a backend that cannot answer, not a small number.
Why does print() inside a jitted JAX function only run once?: Because your Python body only runs once. Every jitted call assembles a key from the argument signature and the abstract input types and looks it up in a cache; on a hit the cached parameters come straight back and the body is never touched, and only a miss traces. So the print fires at trace time, which happens once per distinct key. The key has four moving parts worth knowing. It holds a weak reference to the function object itself, so re-running the cell that defines your function produces a new object and a fresh miss. Static arguments are part of the argument signature. The abstract types carry shape, dtype and sharding, so a new shape misses. And a session-wide tracing context rides on top, so flipping a config anywhere invalidates every entry. To make the print fire again, clear the cache for that function, which evicts the weak reference and then clears the shared parameter cache wholesale. To print on every call instead, use jax.debug.print, which stages a callback into the compiled program; its docstring warns that it does not work with f-strings because formatting is delayed, so the format string and the values are passed separately.
How do I dump the StableHLO or HLO that JAX produces?: Two options, and both must be set before JAX is imported, because both are backed by a holder that reads the environment once at import and has no context manager. JAX_DUMP_IR_TO names a directory, and an empty value means no dumping at all. JAX_DUMP_IR_MODES is a comma-separated list defaulting to stablehlo, and the other three values are jaxpr, jaxpr_html and eqn_count_pprof. Filenames are predictable, which is what makes two runs diffable: modules come out as a four-digit counter, the module name with unsafe characters stripped, and a stage name; jaxprs as a six-digit counter and the function name, with an extension naming the form. The directory is created for you, and the special value sponge redirects everything to a test output directory named by the environment. One detail catches people out: the only site that dumps a module with a stage name sits at the top of the compile entry point, before the persistent cache is consulted, so the file appears even when the compile was served from cache. The before-and-after-pass dumping that XLA does is a separate system driven by XLA flags, and turning it on does not change the compilation cache key, because the dump flags sit on an explicit exclusion list.
Why is my JAX compilation cache directory empty?: Almost always because your compiles are too fast to be worth caching. Three gates sit between a successful compile and a file on disk. The write path times the compile and compares it against a minimum, whose default is one second, and skips anything faster on the reasoning that recompiling is cheaper than reading. A second gate is a minimum entry size in bytes, defaulting to zero but documented as something the runtime may raise to suit the filesystem. A third restricts writes to process zero to avoid contention on shared filesystems, which is free in a notebook but explains sparse caches in multi-host jobs. Every skip logs a line explaining itself, and by default those lines sit at debug level where nobody sees them; setting the cache-miss explanation option raises them to warning level, at which point the empty directory explains itself. Two more things worth separating: the enable switch already defaults to true, so a missing directory setting rather than a disabled cache is the usual configuration error, and clearing the JAX caches does not touch the persistent one, as the docstring of that function says outright.

Writing JAX Notebooks That Show Their Work

Questions

Related explainers