Earlier this year, the PennyLane Team released the first versions of Catalyst, a framework for quantum just-in-time compilation. As our experimental project, Catalyst is a brand new execution pipeline that allows you to scale up your PennyLane workflows by simply decorating them with
@qjit
, enabling rapid prototyping without sacrificing performance.

While we continue working on making Catalyst the best quantum JIT solution out there, we're excited to share with you some important improvements and major new features in the Catalyst v0.3 release โ from the ability to work with native Python control flow, compiler-based backpropagation, quantum control and adjoint support, macOS functionality, and more.
Contents
- Native Python control flow with AutoGraph ๐ฎ
- Backpropagation with Enzyme ๐
- macOS binaries available ๐ฅ๏ธ
- Quantum control and adjoint support ๐ค
- More flexible QNode argument and return handling ๐งโโ๏ธ
- Improvements ๐ ๏ธ
- Breaking changes ๐
- What's next? ๐ฐ๏ธ
- Contributors โ๏ธ
Native Python control flow with AutoGraph ๐ฎ
With Catalyst v0.3, we've introduced AutoGraph, which allows you to write Catalyst-compatible programs using native Python control statements.
This means that Python statements like if
, else
, elif
, and for
can be used with the @qjit
decorator โ Catalyst will automatically capture
loops and conditional statements, and preserve the program structure
under compilation.
dev = qml.device("lightning.qubit", wires=4) @qjit(autograph=True) @qml.qnode(dev) def cost(weights, data): qml.AngleEmbedding(data, wires=range(4)) for x in weights: for j, p in enumerate(x): if p > 0: qml.RX(p, wires=j) elif p < 0: qml.RY(p, wires=j) for j in range(4): qml.CNOT(wires=[j, jnp.mod((j + 1), 4)]) return qml.expval(qml.PauliZ(0) + qml.PauliZ(3))
>>> weights = jnp.linspace(-1, 1, 20).reshape([5, 4]) >>> data = jnp.ones([4]) >>> cost(weights, data) array(0.30455313)
This feature is currently opt-in, it requires TensorFlow to be installed and
can be used by setting autograph=True
in the qjit
decorator.
In some cases there may be a few caveats when using this feature, especially around the use of types and variable assignment. See the AutoGraph documentation for more details.
Backpropagation with Enzyme ๐
Catalyst now supports backpropagation of classical processing in arbitrary hybrid programs, through integration with Enzyme, a tool that differentiates code at the LLVM level.
This allows catalyst.grad
to differentiate hybrid functions that contain
both classical processing and QNodes, via a combination of backpropagation
and quantum gradient methods.
dev = qml.device("lightning.qubit", wires=1) @qml.qnode(dev, diff_method="parameter-shift") def circuit1(x): y = jnp.exp(x[0] ** 2) / jnp.cos(x[1] / 4) qml.RX(y, wires=0) return qml.probs() @qml.qnode(dev, diff_method="adjoint") def circuit2(x): qml.RX(x[0] * x[1] ** 2, wires=0) return qml.expval(qml.PauliZ(0)) def cost(x): return jnp.prod(circuit1(x)) + circuit2(x)
Here, the classical processing will be differentiated using backpropagation, while the QNodes will be differentiated using parameter-shift and adjoint differentiation, respectively:
>>> x = jnp.array([0.4, 0.1]) >>> qjit(grad(cost))(x) array([0.16736644, 0.00098814])
Note that backpropagation support is currently restricted to first-order derivatives
and does not support circuits that involve mid-circuit measurements. However,
you can still use catalyst.grad(func, method="fd")
to compute
finite-difference gradients of any differentiable function.
macOS binaries available ๐ฅ๏ธ
With this release, Catalyst now officially supports macOS ARM devices, such as Apple M1/M2 machines, with macOS binary wheels available on PyPI. Simply run
pip install pennylane-catalyst
to start using Catalyst and @qjit
with PennyLane.
Quantum control and adjoint support ๐ค
Two new functions,
catalyst.ctrl
and
catalyst.adjoint
,
allow for quantum control and the adjoint operation to be represented
and used in compiled functions. Notably, you can use these functions in
conjunction with classical control flow, such as catalyst.cond
and
catalyst.for_loop
.
Turn on AutoGraph, and you can even apply quantum control and adjoint operations on subcircuits that use native Python control flow!
dev = qml.device("lightning.qubit", wires=4) @qjit(autograph=True) @qml.qnode(dev) def circuit(x): qml.RY(1.6, wires=3) def ansatz(): for i in range(3): qml.RX(x / 2, wires=i) catalyst.ctrl(ansatz, control=3)() catalyst.adjoint(ansatz)() return qml.expval(qml.PauliZ(0))
>>> circuit(3.4) array(0.41909689)
In a future release of Catalyst we will merge the behaviour of
catalyst.adjoint
and catalyst.ctrl
into the existing PennyLane functions, qml.adjoint
and qml.ctrl
.
More flexible QNode argument and return handling ๐งโโ๏ธ
There is now a lot more flexibility in both the allowed arguments and the return statement of a QNode.
QJIT-compiled programs now support (nested) container types as inputs and outputs of compiled functions. This includes lists and dictionaries, as well as any data structure implementing the PyTree protocol.
For example, we can now write a program that accepts and returns a mix of dictionaries, arrays, and even operations:
dev = qml.device("lightning.qubit", wires=2) @qjit @qml.qnode(dev) def circuit(params): qml.RX(params["x"][0], wires=0) qml.RY(params["x"][1], wires=1) qml.CNOT(wires=[0, 1]) H = weights["H"] - 0.4 * qml.PauliZ(1) return {"expval": qml.expval(H), "probs[1]": qml.probs(1)}
>>> weights = {"x": jnp.array([0.1, 0.4]), "H": qml.PauliX(0)} >>> circuit(weights) {'expval': array(0.36658381), 'probs[1]': array([0.04177024, 0.95822976])}
Improvements ๐ ๏ธ
In addition to the above new features, we have also been working on a huge number of improvements:
-
JAX arrays can now be indexed and updated via
array.at[index]
. -
A new guide to the current 'sharp bits' of Catalyst, including common patterns that may be useful, is now available in the documentation.
-
The compiler driver has been completely rewritten in C++, improving compile-time performance by avoiding round-tripping, and resulting in significantly improved error handling.
-
Various performance improvements, including:
-
Improved execution and compile times, achieved by generating more efficient code and avoiding unnecessary optimizations.
-
Reduced execution time of compiled functions from the Python frontend, achieved by only loading the user program library once per compilation, avoiding unnecessary array copies and type promotion, and only generating return value types once per compilation.
-
Reduced peak memory utilization of compiled programs, achieved by allowing tensors to be scheduled for deallocation.
-
Breaking changes ๐
As new things are added, outdated features are removed. Here's a summary of what has changed in this release:
-
Due to the change allowing Python container objects as inputs to compiled functions, Python lists are no longer automatically converted to JAX arrays.
This means that indexing on lists when the index is not static will cause a
TracerIntegerConversionError
, consistent with JAX's behaviour:@qjit def f(x: list, index: int): return x[index]
However, if the parameter
x
above is a JAX or NumPy array, the compilation will continue to succeed. -
The
catalyst.grad
functionality for differentiating vector-valued functions has been moved to a separate functioncatalyst.jacobian
, which specifically supports differentiation of functions that return multiple or non-scalar outputs.catalyst.grad
now enforces that it is differentiating a function with a single scalar return value. -
catalyst.grad
has a new default differentiation method,method="auto"
, which will result in backpropagation for classical processing and the QNode-specifieddiff_method
used for computing quantum gradients. -
The JAX version used by Catalyst has been updated to
v0.4.14
; the minimum PennyLane version required is nowv0.32
.
What's next? ๐ฐ๏ธ
In addition to these improvements, the Catalyst team has been hard at work building out the Catalyst infrastructure, to enable us to move on some big new features in the pipeline. This includes:
-
Asynchronous QNode execution,
-
deeper PennyLane integration,
-
support for PennyLane transformations,
-
a device plugin system,
-
hybrid algorithm optimizations,
-
and many others!
These highlights are just scratching the surface โ check out the full release notes for more details.
Contributors โ๏ธ
As always, this release would not have been possible without the hard work of our development team and contributors:
Ali Asadi, David Ittah, Erick Ochoa Lopez, Jacob Mai Peng, Romain Moyard, and Sergei Mironov.
As we continue to build out these features, we encourage you to swing by our GitHub repo and let us know what features you would like to see, help contribute to some of the efforts above, or simply join the discussion.
Down the line, we plan to upstream the Catalyst frontend into PennyLane proper, providing native JIT functionality built into PennyLane. To figure out the details, we need your help โ let us know your use cases by starting a conversation or trying out Catalyst for yourself.
If you are interested in connecting a quantum device with Catalyst, contributing quantum compilation routines, or even adding new frontends, please get in touch with us via GitHub.
In the meantime, make sure to keep an eye on the PennyLane Blog and follow us on Twitter and LinkedIn for the latest Catalyst updates throughout 2023 and beyond.
About the author
Josh Izaac
Josh is a theoretical physicist, software tinkerer, and occasional baker. At Xanadu, he contributes to the development and growth of Xanaduโs open-source quantum software products.