Catalyst v0.2 released

Earlier this year, the PennyLane Team announced (and released!) Catalyst v0.1, 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.

Since the release of Catalyst v0.1 we've been listening to your feedback and chugging away to make improvements and add new features to Catalyst. While Catalyst is still a work-in-progress, we're excited to share with you some important improvements and new features in the Catalyst v0.2 release — and also give you an update on what's to come.

Contents

• Better JAX Integration
• Amazon Braket integration
• New gradient functions
• Improvements
• What's next?
• Contributors

Better JAX Integration

Catalyst programs can now be used inside of a larger JAX workflow which uses JIT compilation, automatic differentiation, and other JAX transforms.

For example, you can call a Catalyst qjit-compiled function from within a JAX jit-compiled function:

import pennylane as qml
from catalyst import qjit
from jax import numpy as jnp
dev = qml.device("lightning.qubit", wires=1)

@qjit
@qml.qnode(dev)
def circuit(x):
    qml.RX(jnp.pi * x[0], wires=0)
    qml.RY(x[1] ** 2, wires=0)
    qml.RX(x[1] * x[2], wires=0)
    return qml.probs(wires=0)

@jax.jit
def cost_fn(weights):
    x = jnp.sin(weights)
    return jnp.sum(jnp.cos(circuit(x)) ** 2)

>>> cost_fn(jnp.array([0.1, 0.2, 0.3]))
Array(1.32269195, dtype=float64)

Catalyst-compiled functions can now also be automatically differentiated to the first order via JAX, both in forward and reverse mode, and vectorized using jax.vmap:

>>> v_cost = jax.vmap(cost_fn)
>>> batched_params = jnp.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]])
>>> jax.jacobian(v_cost)(batched_params)
Array([[[0.49249037, 0.05197949, 0.02991883],
        [0.        , 0.        , 0.        ]],
       [[0.        , 0.        , 0.        ],
        [0.10338042, 0.02040227, 0.01444815]]], dtype=float64)

In particular, this allows for a reduction in boilerplate when using JAX-compatible optimizers such as jaxopt:

>>> opt = jaxopt.GradientDescent(cost_fn)
>>> params = jnp.array([0.1, 0.2, 0.3])
>>> (final_params, _) = jax.jit(opt.run)(params)
>>> final_params
Array([-0.00320799,  0.03475223,  0.29362844], dtype=float64)

Note that, in general, best performance will be seen when the Catalyst @qjit decorator is used to JIT the entire hybrid workflow. However, there may be cases where you might want to delegate only the quantum part of your workflow to Catalyst, and let JAX handle classical components (for example, due to missing a feature or compatibility issue in Catalyst).

Amazon Braket integration

The Amazon Braket integration now enables quantum subprograms within a JIT-compiled Catalyst workflow to execute on Braket simulator and hardware devices, including remote cloud-based simulators such as SV1.

Simply install the PennyLane-Braket plugin and Braket SDK in order to use Braket devices within a qjit-compiled function:

dev = qml.device("braket.local.qubit", backend="braket_sv", wires=2)

@qjit
def workflow(x: float, y: float):

    @qml.qnode(dev)
    def circuit(x, y):
        qml.RX(y * x, wires=0)
        qml.RX(x * 2, wires=1)
        return qml.expval(qml.PauliY(0) @ qml.PauliZ(1))

    return catalyst.grad(circuit)(x, y)

>>> workflow(0.1, 0.2)
array(-0.18802786)

For more details on using PennyLane and Braket, please see the PennyLane-Braket documentation. The full range of available Braket quantum devices and their features can be found in the Amazon Braket Developer Guide.

Internally, the quantum instructions are generating OpenQASM3 kernels at runtime; these are then executed on both local (braket.local.qubit) and remote (braket.aws.qubit) devices backed by Amazon Braket Python SDK, with measurement results then propagated back to the frontend.

Note that not all Catalyst features are supported with Braket at initial release. In particular, dynamic circuit features, such as mid-circuit measurements, will not work with Braket devices.

New gradient functions

New functions catalyst.vjp and catalyst.jvp allow you to directly extract vector-Jacobian and Jacobian-vector products directly within qjit-compiled functions, allowing for more fine-grained and efficient control when building bespoke autodifferentiation pipelines.

@qjit
def vjp(params, cotangent):
    def f(x):
        y = [jnp.sin(x[0]), x[1] ** 2, x[0] * x[1]]
        return jnp.stack(y)

    return catalyst.vjp(f, [params], [cotangent])

>>> x = jnp.array([0.1, 0.2])
>>> dy = jnp.array([-0.5, 0.1, 0.3])
>>> vjp(x, dy)
[array([0.09983342, 0.04      , 0.02      ]),
 array([-0.43750208,  0.07000001])]

Improvements

In addition to the above new features, we have also been working on a huge number of improvements under-the-hood:

• Catalyst conditional functions defined via @catalyst.cond now support an arbitrary number of 'else if' chains, via the new @cond_fn.else_if decorator method.

• Support for multiple backend devices within a single qjit-compiled function is now available.

  For example, if you compile the Catalyst runtime with lightning.kokkos support (via the compilation flag ENABLE_LIGHTNING_KOKKOS=ON), you can use lightning.qubit and lightning.kokkos within a singular workflow.

• Support for returning the variance of Hamiltonians and tensor products via qml.var has been added.

• The catalyst.grad function now supports using the differentiation method defined on the QNode (via the diff_method argument) rather than applying a global differentiation method. This can be enabled by calling catalyst.grad(func, method="defer").

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:

• Bringing the Catalyst autodifferentiation to feature parity with PennyLane, including native, LLVM-layer classical autodifferentiation.

• Support for native Python control flow within qjit-compiled functions;

• Continued improvements to usability and debugging, including the ability to print and draw circuits from within the qjit, support for array assignment, dynamic shaped arrays, and more;

• Availability of additional device support;

• Support for PennyLane transforms;

• Binaries and ease-of-installation for Windows and Mac systems;

• …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.