Compiling workflows

In addition to quantum circuit transformations, PennyLane also supports full hybrid just-in-time (JIT) compilation via the qjit() decorator and various hybrid compilers, which can be installed separately.

The best supported and default compiler is the Catalyst hybrid compiler. Catalyst allows you to compile the entire quantum-classical workflow, including any optimization loops. This maximizes performance and enables running the entire workflow on accelerator devices.

In addition, PennyLane also supports compiling restricted programs via CUDA Quantum; see the CUDA Quantum section below for more details.

Installing compilers

Currently, Catalyst must be installed separately, and only supports the JAX interface and select devices such as lightning.qubit, lightning.kokkos, braket.local.qubit and braket.aws.qubit. It does not support default.qubit.

On MacOS and Linux, Catalyst can be installed with pip:

pip install pennylane-catalyst

Check out the Catalyst documentation for installation instructions.

Just-in-time compilation

Using Catalyst with PennyLane is a simple as using the @qjit decorator to compile your hybrid workflows:

from jax import numpy as jnp

dev = qml.device("lightning.qubit", wires=2, shots=1000)

@qml.qjit
@qml.qnode(dev)
def circuit(params):
    qml.Hadamard(0)
    qml.RX(jnp.sin(params[0]) ** 2, wires=1)
    qml.CRY(params[0], wires=[0, 1])
    qml.RX(jnp.sqrt(params[1]), wires=1)
    return qml.expval(qml.PauliZ(1))

The qjit() decorator can also be used on hybrid functions – that is, functions that include both QNodes and classical processing.

@qml.qjit
def hybrid_function(params, x):
    grad = qml.grad(circuit)(params)
    return jnp.abs(grad - x) ** 2

In addition, functions that are compiled with @jax.jit can contain calls to qjit-compiled functions. For example, below we compile a full optimization loop, using @jax.jit:

import jaxopt

@jax.jit
def optimization():
    # initial parameter
    params = jnp.array([0.54, 0.3154])

    # define the optimizer
    opt = jaxopt.GradientDescent(circuit, stepsize=0.4)
    update = lambda i, args: tuple(opt.update(*args))

    # perform optimization loop
    state = opt.init_state(params)
    (params, _) = jax.lax.fori_loop(0, 100, update, (params, state))

    return params

Compiling the entire hybrid workflow using @qml.qjit however will lead to better performance. For more details, please see the Catalyst documentation.

Control flow

The Catalyst compiler also supports capturing imperative Python control flow in compiled programs, resulting in control flow being interpreted at runtime rather than in Python at compile time. You can enable this feature via the autograph=True keyword argument.

@qml.qjit(autograph=True)
@qml.qnode(dev)
def circuit(x: int):

    if x < 5:
        qml.Hadamard(wires=0)
    else:
        qml.T(wires=0)

    return qml.expval(qml.PauliZ(0))

>>> circuit(3)
array(0.)
>>> circuit(5)
array(1.)

Note that AutoGraph results in additional restrictions, in particular whenever global state is involved. Please refer to the AutoGraph guide for a complete discussion of the supported and unsupported use-cases.

CUDA Quantum

The PennyLane qjit() decorator can also be used to compile programs using CUDA Quantum, a hybrid compiler toolchain by NVIDIA.

First, Catalyst and CUDA Quantum need to be installed:

pip install pennylane-catalyst cuda_quantum

Then, simply specify compiler="cuda_quantum" in the @qjit decorator:

dev = qml.device("softwareq.qpp", wires=2)

@qml.qjit(compiler="cuda_quantum")
@qml.qnode(dev)
def circuit(x):
    qml.RX(x[0], wires=0)
    qml.RY(x[1], wires=1)
    qml.CNOT(wires=[0, 1])
    return qml.expval(qml.PauliY(0))

>>> circuit(jnp.array([0.5, 1.4]))
-0.47244976756708373

The following devices are available when compiling with CUDA Quantum:

• softwareq.qpp: a modern C++ statevector simulator

• nvidia.custatevec: The NVIDIA CuStateVec GPU simulator (with support for multi-gpu)

• nvidia.cutensornet: The NVIDIA CuTensorNet GPU simulator (with support for matrix product state)

Note that CUDA Quantum compilation currently does not have feature parity with Catalyst compilation; in particular, AutoGraph, control flow, differentiation, and various measurement statistics (such as probabilities and variance) are not yet supported.

Additional resources

For more details on using the qjit() decorator and Catalyst with PennyLane, please refer to the Catalyst quickstart guide, as well as the sharp bits and debugging tips page for an overview of the differences between Catalyst and PennyLane, and how to best structure your workflows to improve performance when using Catalyst.

To make your own compiler compatible with PennyLane, please see the compiler module documentation.

introduction/compiling_workflows

 

Download Python script

 

Download Notebook

 

View on GitHub