Catalyst v0.3 released

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 function catalyst.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-specified diff_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 now v0.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.