El equipo de PennyLane, una biblioteca de software cuántico, estaba desarrollando Catalyst, un compilador cuántico basado en MLIR, cuando descubrió que su infraestructura podía ejecutar código JAX clásico sin necesidad de instrucciones cuánticas. El decorador @qjit acepta funciones JAX puras con NumPy y control de flujo nativo de Python, elude por completo el backend de XLA y compila directamente hasta código máquina mediante LLVM, con soporte para retropropagación a través de Enzyme.
El proyecto nació con tres objetivos: aprovechar la cadena de herramientas clásica madura de LLVM en lugar de reinventarla, preservar la estructura de los programas cuánticos híbridos con bucles y condicionales, y mantener la autodiferenciación cuántica. Para ello, el equipo extendió JAX con conjuntos de instrucciones cuánticas, captura de control de flujo en Python y arrays de forma dinámica. Tras interceptar la representación StableHLO que JAX genera, Catalyst la convierte a dialectos MLIR de propósito general (linalg, arith, scf) y la envía a LLVM y Enzyme, prescindiendo del runtime de XLA.
Los propios autores reconocen que este enfoque no pretende competir con XLA en cargas de aprendizaje profundo estándar sobre GPU o TPU, donde XLA sigue siendo superior. Sin embargo, abrir JAX al ecosistema LLVM abre posibilidades como generar binarios Ahead-of-Time autónomos sin dependencias de ML, soportar arrays de forma dinámica y simplificar la compilación al eliminar la dependencia de Bazel para construir XLA. El artículo describe Catalyst como un puente para experimentar entre JAX y LLVM, más allá del compilador cuántico original.