How does JAX turn a jaxpr into StableHLO?

By walking the jaxpr equation by equation and calling a registered lowering rule for each primitive. The walk is jaxpr_subcomp in jax/_src/interpreters/mlir.py: it keeps a dictionary mapping jaxpr variables to MLIR values, reads each equation’s inputs out of it, looks up a rule for the equation’s primitive, and writes the results back. A rule is an ordinary Python function that takes a context object plus the input MLIR values, emits StableHLO operations at the current insertion point, and returns handles to their results. The loop is wrapped by lower_jaxpr_to_module, which builds the module, sets the replica and partition attributes, builds the main function with the runtime’s argument conventions, verifies the result, and hands it back. Lowering does no scheduling and no optimisation: the operations come out in the order the equations appear, and every rearrangement you see later was done by XLA.

How many lowering rules does JAX have, and how do I find the rule for a primitive?

At the pinned commit there are 380 direct calls to mlir.register_lowering in the jax tree outside tests, plus eleven calls to a linear-algebra helper, each of which expands into two to four registrations. Resolving those and expanding the legacy gpu alias gives 462 registrations across 311 distinct primitives in 65 files. To find one, grep the tree for register_lowering with the primitive name: the rule is almost always registered on the line right after the primitive is defined, in the same file. 243 of the 311 primitives have exactly one rule that applies on every platform, 68 have at least one per-platform variant, and 26 have no platform-agnostic rule at all.

What does the error "MLIR translation rule for primitive not found for platform" mean?

It means the lowering-rule lookup found nothing for that primitive on the platform you are lowering for, and there was no platform-agnostic rule to fall back to. It is raised by lower_per_platform, in Python, after tracing has already succeeded and before any MLIR module is built. Twenty-six primitives at the pin have no shared rule, and they are mostly the ones whose rules emit a call into a vendor library that only exists for some backends. The fix is either a rule for your platform or a different primitive; JAX deliberately does not fall back to a host implementation, because silently copying arrays to the host inside a compiled program would introduce a synchronising round trip nobody asked for.

Why does jnp.linalg.eig fail on TPU when jnp.linalg.eigh works?

Because of what is in the registry. The eig primitive has four lowering registrations at the pin: one for cpu and three for cuda, rocm and oneapi, all of them emitting a call into a vendor eigendecomposition routine. There is no TPU entry and no platform-agnostic rule, so lowering raises NotImplementedError naming the primitive and the platform. The symmetric decomposition eigh does have a TPU rule, registered from a JAX-side implementation, which is why the two behave differently on the same machine. If you need the general case on a TPU, compute it as a host step and treat the transfer as deliberate.

Where do sharding annotations live in the StableHLO module JAX emits?

On the arguments and results of the entry function, as attributes, not as operations in the body. lower_jaxpr_to_fun attaches sdy.sharding when the Shardy partitioner is on, which is the default at the pinned commit, and mhlo.sharding when it is off. The same attribute dictionary carries mhlo.memory_kind for the memory space, mhlo.layout_mode for a concrete layout or the literal string auto, and donation as either tf.aliasing_output when JAX matched the donated input to an output or jax.buffer_donor when it left the decision to the compiler. The partitioner sees only those annotations plus the shapes, which is why every collective in the compiled program was either written by you or inserted downstream. The one exception is a shard_map body, which lowers to an explicit manual-computation operation inside the module.

Why does a lowered module have more arguments than my function?

Because three groups of arguments are added ahead of yours, always in the same order: dimension variables first, then one token per ordered effect, then any constants hoisted out of the jaxpr body, then your own arguments. Each group is marked with an attribute so it can be told apart: jax.global_constant carrying the variable name for a dimension variable, jax.token for a token, jax.const for a hoisted constant. The return side follows the same convention, with token results first. Dimension variables appear when a shape is polymorphic, which is how one exported module serves any batch size; a rule that needs a real size at lowering time evaluates the symbolic expression against those argument values.

From Jaxpr to StableHLO: One Rule Per Primitive

The last thing JAX does before handing your program to XLA is walk the jaxpr and, for every equation in it, look up a function in a dictionary keyed by the primitive. This page opens that dictionary: 311 primitives and 462 registrations counted from the call sites at one pinned commit of jax-ml/jax, filterable by platform, every row a link to the line that registers it. Then the wrapper around the loop, in the order the code builds it: the arguments you never passed and why they are in that order, the sharding and layout and donation attributes the partitioner actually reads, the tokens that give two debug prints an order, and the error you get when nobody wrote a rule for your accelerator.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of From Jaxpr to StableHLO: One Rule Per Primitive — built to be understood, not skimmed.

Questions

How does JAX turn a jaxpr into StableHLO?: By walking the jaxpr equation by equation and calling a registered lowering rule for each primitive. The walk is jaxpr_subcomp in jax/_src/interpreters/mlir.py: it keeps a dictionary mapping jaxpr variables to MLIR values, reads each equation’s inputs out of it, looks up a rule for the equation’s primitive, and writes the results back. A rule is an ordinary Python function that takes a context object plus the input MLIR values, emits StableHLO operations at the current insertion point, and returns handles to their results. The loop is wrapped by lower_jaxpr_to_module, which builds the module, sets the replica and partition attributes, builds the main function with the runtime’s argument conventions, verifies the result, and hands it back. Lowering does no scheduling and no optimisation: the operations come out in the order the equations appear, and every rearrangement you see later was done by XLA.
How many lowering rules does JAX have, and how do I find the rule for a primitive?: At the pinned commit there are 380 direct calls to mlir.register_lowering in the jax tree outside tests, plus eleven calls to a linear-algebra helper, each of which expands into two to four registrations. Resolving those and expanding the legacy gpu alias gives 462 registrations across 311 distinct primitives in 65 files. To find one, grep the tree for register_lowering with the primitive name: the rule is almost always registered on the line right after the primitive is defined, in the same file. 243 of the 311 primitives have exactly one rule that applies on every platform, 68 have at least one per-platform variant, and 26 have no platform-agnostic rule at all.
What does the error "MLIR translation rule for primitive not found for platform" mean?: It means the lowering-rule lookup found nothing for that primitive on the platform you are lowering for, and there was no platform-agnostic rule to fall back to. It is raised by lower_per_platform, in Python, after tracing has already succeeded and before any MLIR module is built. Twenty-six primitives at the pin have no shared rule, and they are mostly the ones whose rules emit a call into a vendor library that only exists for some backends. The fix is either a rule for your platform or a different primitive; JAX deliberately does not fall back to a host implementation, because silently copying arrays to the host inside a compiled program would introduce a synchronising round trip nobody asked for.
Why does jnp.linalg.eig fail on TPU when jnp.linalg.eigh works?: Because of what is in the registry. The eig primitive has four lowering registrations at the pin: one for cpu and three for cuda, rocm and oneapi, all of them emitting a call into a vendor eigendecomposition routine. There is no TPU entry and no platform-agnostic rule, so lowering raises NotImplementedError naming the primitive and the platform. The symmetric decomposition eigh does have a TPU rule, registered from a JAX-side implementation, which is why the two behave differently on the same machine. If you need the general case on a TPU, compute it as a host step and treat the transfer as deliberate.
Where do sharding annotations live in the StableHLO module JAX emits?: On the arguments and results of the entry function, as attributes, not as operations in the body. lower_jaxpr_to_fun attaches sdy.sharding when the Shardy partitioner is on, which is the default at the pinned commit, and mhlo.sharding when it is off. The same attribute dictionary carries mhlo.memory_kind for the memory space, mhlo.layout_mode for a concrete layout or the literal string auto, and donation as either tf.aliasing_output when JAX matched the donated input to an output or jax.buffer_donor when it left the decision to the compiler. The partitioner sees only those annotations plus the shapes, which is why every collective in the compiled program was either written by you or inserted downstream. The one exception is a shard_map body, which lowers to an explicit manual-computation operation inside the module.
Why does a lowered module have more arguments than my function?: Because three groups of arguments are added ahead of yours, always in the same order: dimension variables first, then one token per ordered effect, then any constants hoisted out of the jaxpr body, then your own arguments. Each group is marked with an attribute so it can be told apart: jax.global_constant carrying the variable name for a dimension variable, jax.token for a token, jax.const for a hoisted constant. The return side follows the same convention, with token results first. Dimension variables appear when a shape is polymorphic, which is how one exported module serves any batch size; a rule that needs a real size at lowering time evaluates the symbolic expression against those argument values.

From Jaxpr to StableHLO: One Rule Per Primitive

Questions

Related explainers