The Catalyst team finds it built a general JAX-to-LLVM compiler by accident
A quantum computing team building Catalyst, a compiler for the PennyLane quantum software library, set out to solve a narrower problem: they needed a fast way to capture classical Python processing (NumPy and SciPy code) and represent it in MLIR, the intermediate representation their quantum pipeline is built on. They chose JAX because it can trace Python functions into a computational graph, covers most of the NumPy and SciPy APIs, and already lowers to MLIR. On top of that they added native Python control flow support (via a fork of TensorFlow's Autograph, which they renamed Diastatic Malt) and dynamically shaped arrays, both needed because quantum algorithms often involve loops and if-statements whose shape depends on measurement outcomes.
While wiring JAX into their quantum pipeline, the team realized their @qjit decorator does not require any quantum instructions at all. Feed it pure JAX NumPy code with ordinary Python control flow, and Catalyst bypasses Google's XLA compiler backend completely: it takes the JAX representation, lowers it directly into standard, general-purpose MLIR dialects such as linalg, arith and scf, and hands that straight to LLVM to produce machine code, using the Enzyme tool for automatic differentiation and backpropagation. XLA itself does eventually call LLVM too, but only after routing everything through its own StableHLO dialect and a heavy JIT runtime (PJRT); Catalyst's shortcut drops that runtime and its optimizations altogether.
The authors are explicit that this will not beat XLA for standard deep learning workloads: "this is not going to beat XLA for standard deep learning workloads," and anyone training a large transformer should stick with plain JAX. What the shortcut does enable, in their view, is a set of side benefits from working at the LLVM level directly: standalone ahead-of-time binaries that need no Python or ML runtime to execute; native support for dynamically shaped arrays without XLA's usual recompilation penalty when input dimensions change; full native Python control flow and NumPy-style indexing, all differentiable; and the ability to target bare-metal or edge hardware, such as FPGAs or custom embedded chips, that cannot carry a full ML runtime. Because the code hits standard MLIR before reaching hardware, it is also straightforward to inject custom MLIR dialects and compiler passes into the pipeline.
The team says it built this classical JAX-to-LLVM path purely as a side effect of solving their quantum problem, and that they are not certain yet what it is good for beyond their own use case: "Honestly? We aren't entirely sure yet." The post asks readers working in custom hardware, edge machine learning or compiler infrastructure whether dropping the XLA runtime in favor of a direct JAX-to-MLIR-to-LLVM path solves a problem for them. No benchmark numbers, release dates or maturity claims for the bypass path are given; the package is installable via pip install pennylane-catalyst.
Key facts
- Catalyst, a compiler for the PennyLane quantum library, traces classical JAX/NumPy code into MLIR and found it can compile pure non-quantum JAX code with zero quantum instructions
- The path bypasses Google's XLA backend entirely, lowering JAX straight into standard MLIR dialects (linalg, arith, scf) and then to machine code via LLVM, with Enzyme handling backprop
- Benefits cited: standalone ahead-of-time binaries with no ML runtime dependency, dynamically shaped arrays without XLA's recompilation penalty, native differentiable Python control flow, and edge or bare-metal deployment on FPGAs and embedded chips
- The authors say this will not beat XLA for standard deep learning workloads and that anyone training a large transformer should stick with plain JAX
- The team, focused on quantum computing, says it is not yet sure of the broader non-quantum usefulness of the pipeline and is asking readers in edge ML and compiler infrastructure for feedback
Why it matters
Most JAX code reaches machine code through XLA, a heavily optimized but fairly closed pipeline tied to a JIT runtime and to GPU/TPU targets. This project shows that JAX's own tracing and MLIR lowering can be repurposed to skip XLA altogether and go straight to standard LLVM infrastructure, which is exactly the layer that custom hardware, edge devices and compiler tooling already know how to work with. It was found by accident, while solving an unrelated quantum computing problem, which is itself notable: a niche compiler project for quantum-classical hybrid workflows turned up a general-purpose classical compilation path nobody had set out to build.
Who it affects
Most directly this affects the Catalyst and PennyLane quantum computing project itself, which needed the classical compilation layer to support hybrid quantum-classical workflows and gradients. Beyond that, the authors are explicitly asking three groups whether it is useful to them: people building custom hardware or accelerators, people deploying machine learning models to edge devices such as FPGAs or embedded chips, and people working on compiler infrastructure who might want to inject their own MLIR dialects and passes. Standard deep learning practitioners training large models are told directly that this path is not for them.
How to use it
The package is installed with pip install pennylane-catalyst. A workflow is wrapped in the @qjit decorator with autograph=True to enable native Python control flow (while loops, if statements) without needing jax.lax.while_loop or jax.lax.cond; ordinary jax.numpy code and NumPy-style array operations work inside it. No quantum instructions are required for the classical-only path to trigger. No pricing applies since the tool is an open-source library; no version number, release date or stability guarantee for this specific bypass behavior is given in the source.
How solid is it
The account comes directly from the team that built Catalyst, describing their own pipeline and their own discovery, which makes the technical description first-hand. However, the source gives no benchmark numbers, no timing comparisons against XLA, and no data on how mature or production-tested this classical bypass path is; the authors themselves say they are not certain of its broader usefulness. It should be read as an early, informal engineering finding rather than a validated or benchmarked result.
Risks and caveats
The authors are explicit that this path will not beat XLA for standard deep learning workloads and that large transformer training should stay on plain JAX, so the pipeline is not a general XLA replacement. No performance figures, stability guarantees, or evidence of production use are provided, and the team itself frames the classical-only capability as a side effect they have not fully evaluated rather than a finished feature. No individual author, team name beyond the project, or specific employer is stated in the text itself, only an implied link to Xanadu through a careers-page reference.
“Honestly? We aren't entirely sure yet.”
— the Catalyst team