Sometimes when you build software, you shave a yak so deeply that you end up with a surprisingly nice sweater (while the yak, presumably, wonders why it is suddenly rather drafty).

For the past while, my team and I have been building a quantum compiler called Catalyst for the quantum software library PennyLane . Our goal was relatively straightforward: optimize large, hybrid quantum-classical workflows in a scalable manner using MLIR . To do this, we needed a fast, robust way to capture classical Python processing (including NumPy and associated scientific libraries) and represent it in our intermediate representation.

We chose JAX , for a couple of reasons: its ability to trace through Python functions and capture the computational graph 1 , its support of the NumPy and SciPy APIs with relatively good coverage, and the fact that it already lowers down to MLIR .…