qml.compiler

This module provides support for hybrid quantum-classical compilation. Through the use of the qjit() decorator, entire workflows can be just-in-time (JIT) compiled — including both quantum and classical processing — down to a machine binary on first function execution. Subsequent calls to the compiled function will execute the previously-compiled binary, resulting in significant performance improvements.

Currently, PennyLane supports the Catalyst hybrid compiler and the CUDA Quantum compiler with the qjit() decorator. A significant benefit of Catalyst is the ability to preserve complex control flow around quantum operations — such as if statements and for loops, and including measurement feedback — during compilation, while continuing to support end-to-end autodifferentiation.

Note

Catalyst currently only supports the JAX interface of PennyLane.

Overview

The main entry point to hybrid compilation in PennyLane is via the qjit() decorator that can be mixed with other compiler-specific decorators and functions:

qjit([fn, compiler])	A decorator for just-in-time compilation of hybrid quantum programs in PennyLane.
for_loop([start, ]stop[, step])	A qjit() compatible for-loop for PennyLane programs.
while_loop(cond_fn)	A qjit() compatible for-loop for PennyLane programs.
jvp(f, params, tangents[, method, h, argnum])	A qjit() compatible Jacobian-vector product of PennyLane programs.
vjp(f, params, cotangents[, method, h, argnum])	A qjit() compatible Vector-Jacobian product of PennyLane programs.

In addition, several developer functions are available to probe available hybrid compilers.

available_compilers()	Load and return a list of available compilers that are installed and compatible with the qjit() decorator.
available([compiler])	Check the availability of the given compiler package.
active_compiler()	Check which compiler is activated inside a qjit() evaluation context.
active()	Check whether the caller is inside a qjit() evaluation context.

Presented below is the list of qjit() compatible PennyLane primitives.

adjoint(fn[, lazy])	Create the adjoint of an Operator or a function that applies the adjoint of the provided function.
cond(condition[, true_fn, false_fn, elifs])	Quantum-compatible if-else conditionals --- condition quantum operations on parameters such as the results of mid-circuit qubit measurements.
ctrl(op, control[, control_values, work_wires])	Create a method that applies a controlled version of the provided op.
grad(func[, argnum, method, h])	Returns the gradient as a callable function of hybrid quantum-classical functions.
jacobian(func[, argnum, method, h])	Returns the Jacobian as a callable function of vector-valued (functions of) QNodes.

Compiler

The compiler module provides the infrastructure to integrate external hybrid quantum-classical compilers with PennyLane, but does not provide a built-in compiler.

Currently, only the Catalyst hybrid compiler and CUDA Quantum compiler toolchains are supported with PennyLane, however there are plans to incorporate additional compilers in the near future.

Note

To install Catalyst, simply run the following pip command:

pip install pennylane-catalyst

See the installation guide for more information and supported platforms.

Basic usage

Note

Supported backend devices for Catalyst include lightning.qubit, lightning.kokkos, lightning.gpu, and braket.aws.qubit, but not default.qubit.

For a full list of supported devices, please see Supported devices.

When using just-in-time (JIT) compilation, the compilation is triggered at the call site the first time the quantum function is executed. For example, circuit is compiled in the first call.

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

@qml.qjit
@qml.qnode(dev)
def circuit(theta):
    qml.Hadamard(wires=0)
    qml.RX(theta, wires=1)
    qml.CNOT(wires=[0,1])
    return qml.expval(qml.Z(1))

>>> circuit(0.5)  # the first call, compilation occurs here
array(0.)
>>> circuit(0.5)  # the precompiled quantum function is called
array(0.)

Alternatively, if argument type hints are provided, compilation can occur ‘ahead of time’ when the function is decorated.

from jax.core import ShapedArray

@qml.qjit  # compilation happens at definition
@qml.qnode(dev)
def circuit(x: complex, z: ShapedArray(shape=(3,), dtype=jnp.float64)):
    theta = jnp.abs(x)
    qml.RY(theta, wires=0)
    qml.Rot(z[0], z[1], z[2], wires=0)
    return qml.state()

>>> circuit(0.2j, jnp.array([0.3, 0.6, 0.9]))  # calls precompiled function
array([0.75634905-0.52801002j, 0. +0.j,
    0.35962678+0.14074839j, 0. +0.j])

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.Z(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.

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.

Adding a compiler

Warning

The PennyLane compiler API is experimental and subject to change.

To register any compiler packages, an experimental interface is available. This interface exposes the entry_points metadata under the designated group name pennylane.compilers, including the following entry points:

• compiler_name.context": Path to the compilation evaluation context manager. This context manager should have the method context.is_tracing(), which returns True if called within a program that is being traced or captured.

• compiler_name.ops: Path to the compiler operations module. This operations module may contain compiler specific versions of PennyLane operations, for example cond(), measure(), and adjoint(). Within a JIT context, PennyLane operations may dispatch to these functions.

• compiler_name.qjit: Path to the JIT decorator provided by the compiler. This decorator should have the signature qjit(fn, *args, **kwargs), where fn is the function to be compiled.

where compiler_name should be replaced with the name of the compiler. For example, for Catalyst, we define the entry points catalyst.context, catalyst.ops and catalyst.qjit. This allows the catalyst package to define multiple compilers.

The name of the compiler can then be used by the user to denote which compiler should be used. For example:

@qml.qjit(compiler="catalyst")
def function(x, y):
    ...

@qml.qjit(compiler="compiler_name")
def function(x, y):
    ...

In order to support applying the qjit decorator with and without arguments,

@qml.qjit
def function(x, y):
    ...

@qml.qjit(verbose=True, additional_args, ...)
def function(x, y):
    ...

You should ensure that the qjit decorator itself returns a decorator if no function is provided:

def qjit(fn=None, **kwargs):
    if fn is not None:
        return compile_fn(fn, **kwargs)

    def wrapper_fn(fn):
        return compile_fn(fn, **kwargs)

    return wrapper_fn

code/qml_compiler

 

Download Python script

 

Download Notebook

 

View on GitHub