How to write quantum function transforms in PennyLane

The latest release of PennyLane, version 0.17, contains a number of new features based on differentiable transforms. Transforms can modify the behaviour of quantum operations and functions, or return new functions that serve to process their outputs. They have an increasingly large number of use cases in PennyLane, such as quantum circuit compilation, making an operation controlled or taking its adjoint, as well as computing gradients. Differentiability is a core property of transforms in PennyLane—not only do they preserve the differentiability of quantum functions, they can themselves be differentiated, trained, and optimized. Furthermore, transforms can composed to create fully differentiable pipelines.

In this guide, we provide step-by-step instructions to create two types of differentiable transforms using the features available in PennyLane’s transforms module:

• Tape transforms
• Quantum function transforms

Both act on quantum tapes, which are the low-level data structure used in PennyLane to record and execute quantum operations. As a user or developer, you are likely more used to working at the level of quantum functions and QNodes. Both use tapes under the hood; tape and quantum function transforms involve manipulating and modifying parts of a tape in order to change its effects.

Tape transforms

Tape transforms work exactly as their name suggests: they take one (or more) tapes as input, and return one (or more) tapes as output. In this guide, we focus on the most basic case where both the input and output are a single tape.

As a first example, consider the task of rewriting gates in terms of other gates (transpilation). Suppose we have a circuit containing CNOTs, but our quantum processor can only implement CZs. There is a circuit identity relating the two (a CZ with a Hadamard on the target qubit on each side), but we need a way to apply said identity at every instance of a CNOT for arbitrary circuits.

Let’s start with the tape below. Note that normally, the explicit construction of a tape like this happens within a QNode when it is executed, so we are showing it here mainly for expository purposes.

import pennylane as qml
from pennylane import numpy as np

with qml.tape.QuantumTape() as tape:
    # Some operations
    qml.RX(0.3, wires=0)
    qml.CNOT(wires=[0, 1])
    qml.RY(0.4, wires=1)
    qml.CNOT(wires=[1, 0])

    # Some measurements
    qml.expval(qml.PauliZ(0))
    qml.expval(qml.PauliY(1))

There are two ways to approach constructing a tape transform. One is to write a function that explicitly accepts a tape, modifies it, and returns a new tape. The requires manual construction of the tapes and more in-depth management of the queuing contexts in which the operations are modified and applied. A more convenient way is to use the built-in single_tape_transform decorator. The construction of the original tape and new tape is all handled by the decorator; it also manages the return of the modified tape. Using the decorator, the only part of the transform we need to write is the part that manipulates the operations themselves.

@qml.single_tape_transform
def convert_cnots(tape):
    # Loop through all items in the original tape
    for op in tape.operations + tape.measurements:

        # If it's a CNOT, replace it using the circuit identity
        if op.name == "CNOT":
            wires = op.wires
            qml.Hadamard(wires=wires[1])
            qml.CZ(wires=[wires[1], wires[0]])
            qml.Hadamard(wires=wires[1])

        # If it's not a CNOT, keep the operation as-is and apply it
        else:
            qml.apply(op)

This transform can now be applied directly to a tape, returning a new, transformed tape:

>>> converted_tape = convert_cnots(tape)
>>> converted_tape.operations
[RX(0.3, wires=[0]),
 Hadamard(wires=[1]),
 CZ(wires=[1, 0]),
 Hadamard(wires=[1]),
 RY(0.4, wires=[1]),
 Hadamard(wires=[0]),
 CZ(wires=[0, 1]),
 Hadamard(wires=[0])]

Every CNOT on the original tape was correctly replaced with Hadamard and CZ!

Transforms are composable—we can chain multiple transforms together, which enables us to keep the transforms themselves focused on individual tasks, as well as manage the order in which transforms are performed. For example, we could write a second transform, expand_hadamards, that converts each Hadamard to its decomposition [RZ(np.pi), RY(np.pi/2)] (up to a global phase).

@qml.single_tape_transform
def expand_hadamards(tape):
    for op in tape.operations + tape.measurements:
        if op.name == "Hadamard":
            qml.RZ(np.pi, wires=op.wires)
            qml.RY(np.pi / 2, wires=op.wires)
        else:
            qml.apply(op)

We can create a pipeline of transforms and run the tape through each one in sequence:

>>> pipeline = [convert_cnots, expand_hadamards]
>>> transformed_tape = tape
>>> for transform in pipeline:
...     transformed_tape = transform(transformed_tape)
...
>>> transformed_tape.operations
[RX(0.3, wires=[0]),
 RZ(3.141592653589793, wires=[1]),
 RY(1.5707963267948966, wires=[1]),
 CZ(wires=[1, 0]),
 RZ(3.141592653589793, wires=[1]),
 RY(1.5707963267948966, wires=[1]),
 RY(0.4, wires=[1]),
 RZ(3.141592653589793, wires=[0]),
 RY(1.5707963267948966, wires=[0]),
 CZ(wires=[0, 1]),
 RZ(3.141592653589793, wires=[0]),
 RY(1.5707963267948966, wires=[0])]

Tape transforms are flexible and reusable functions that change the behaviour of a quantum circuit. However, programming in PennyLane is typically done at the level of quantum functions, rather than tapes. Tape transforms can be leveraged to build more convenient quantum function transforms.

Quantum function transforms

In PennyLane, quantum circuits are represented by quantum functions. These are simply Python functions that perform one or more quantum operations, may do classical processing, and return a measurement, all the while remaining differentiable. Quantum functions comprise one part of a QNode (the other being a device). When a QNode is created, a quantum function is bound to it, and the operations of that function are used to create the internal tape that is executed on the device.

Quantum function transforms allow us to work at the level of operations, or entire functions, all before they are used within a QNode. Writing quantum function transforms is straightforward once you know how to write tape transforms. We can construct a quantum function transform by simply adding the qml.qfunc_transform decorator to what would otherwise be a tape transform. The resulting transform takes a quantum function as input, and produces a quantum function as output. The transforms may also have additional parameters. The presence of parameters leads to a slightly different syntax, so we cover both cases below.

Transforms without parameters

Let’s revisit our CNOT conversion tape transform. Instead of the single_tape_transform decorator, we use the qfunc_transform decorator:

@qml.qfunc_transform
def convert_cnots(tape):
    for op in tape.operations + tape.measurements:
        if op.name == "CNOT":
            wires = op.wires
            qml.Hadamard(wires=wires[1])
            qml.CZ(wires=[wires[1], wires[0]])
            qml.Hadamard(wires=wires[1])
        else:
            qml.apply(op)

Notice how, despite being a quantum function transform, the input is still a tape. The decorator does the heavy lifting here and “elevates” the tape transform to a function transform, so that we can pass it quantum functions as inputs. For example, suppose we have the function:

def circuit(x, y):
    qml.RX(x, wires=0)
    qml.RY(y, wires=1)
    qml.CNOT(wires=[0, 1])
    return qml.expval(qml.PauliZ(0))

One way to apply the convert_cnots transform is to apply it directly to the function, producing a new function. That transformed function can then be used to construct a QNode.

>>> transformed_circuit = convert_cnots(circuit)
>>> dev = qml.device('default.qubit', wires=2)
>>> qnode = qml.QNode(transformed_circuit, dev)
>>> qnode(0.3, 0.4)
tensor([0.95533649], requires_grad=True)
>>> print(qml.draw(qnode)(0.3, 0.4))
0: ──RX(0.3)─────╭Z─────┤ ⟨Z⟩
1: ──RY(0.4)──H──╰C──H──┤

Additional flexibility lies in the fact that quantum function transforms themselves can also be used as decorators. For example:

@convert_cnots
def different_circuit(x):
    qml.RX(x, wires=0)
    qml.CNOT(wires=[0, 1])
    qml.CNOT(wires=[1, 2])
    qml.Hadamard(wires=2)
    return qml.expval(qml.PauliZ(0))

When we run this function we see that the transform is automatically applied!

>>> dev = qml.device('default.qubit', wires=3)
>>> qnode = qml.QNode(different_circuit, dev)
>>> qnode(0.2)
tensor([0.98006658], requires_grad=True)
>>> print(qml.draw(qnode)(0.2))
 0: ──RX(0.2)──╭Z───────────────┤ ⟨Z⟩
 1: ──H────────╰C──H──╭Z────────┤
 2: ──H───────────────╰C──H──H──┤

Note that even though the signatures of circuit and different_circuit differ, we are able to apply the same quantum transform on both with ease. This becomes even more powerful as you can apply multiple decorators and transform your function without having to deal with the underlying tapes (provided the tape transforms are already available). To exemplify this, let’s also rewrite our expand_hadamards tape transform as a quantum function transform:

@qml.qfunc_transform
def expand_hadamards(tape):
    for op in tape.operations + tape.measurements:
        if op.name == "Hadamard":
            qml.RZ(np.pi, wires=op.wires)
            qml.RY(np.pi / 2, wires=op.wires)
        else:
            qml.apply(op)

We can create a circuit with both decorators:

@expand_hadamards
@convert_cnots
def different_circuit(x):
    qml.RX(x, wires=0)
    qml.CNOT(wires=[0, 1])
    qml.CNOT(wires=[1, 2])
    qml.Hadamard(wires=2)
    return qml.expval(qml.PauliZ(0))

Let’s run it:

>>> qnode = qml.QNode(different_circuit, dev)
>>> qnode(0.2)
tensor([0.98006658], requires_grad=True)
>>> print(qml.draw(qnode)(0.2))
 0: ──RX(0.2)─────────────╭Z──────────────────────────────────────────────────────────────────┤ ⟨Z⟩
 1: ──RZ(3.14)──RY(1.57)──╰C──RZ(3.14)──RY(1.57)──╭Z──────────────────────────────────────────┤
 2: ──RZ(3.14)──RY(1.57)──────────────────────────╰C──RZ(3.14)──RY(1.57)──RZ(3.14)──RY(1.57)──┤

We get the circuit we expect. An important thing to note is that the order of the decorators affects the order in which the transforms are applied. The order goes from inside (closest to the function definition) to outside (furthest), so in this case, convert_cnots is applied before expand_hadamards, which is why even the Hadamards that occurred from the CNOT conversion are expanded.

Transforms with parameters

Quantum function transforms can also include parameters that will affect their structure and behaviour. For example, CNOT pair insertion, which is commonly used in error-mitigation pipelines, can be expressed as a quantum function transform with an argument that specifies the number of added CNOTs.

As a concrete example, suppose we want to simulate a noisy quantum system in which every RX operation is slightly over-rotated by a small angle. We can write a quantum function transform:

@qml.qfunc_transform
def overrotate_rx(tape, overrot_angle=0):
    for op in tape.operations + tape.measurements:
        if op.name == "RX":
            qml.RX(op.parameters[0] + overrot_angle, wires=op.wires)
        else:
            qml.apply(op)

We can now apply this to an arbitrary quantum function in two ways. One is the decorator method, in which we must specify a particular over-rotation angle when defining the circuit:

@overrotate_rx(0.05)
def circuit(x):
    qml.RX(x, wires=0)
    qml.CNOT(wires=[0, 1])
    qml.CNOT(wires=[1, 2])
    qml.Hadamard(wires=2)
    return qml.expval(qml.PauliZ(0))

If we run this circuit, we find that angles of RX gates are always increased by 0.05:

>>> qnode = qml.QNode(circuit, dev)
>>> qnode(0.3)
tensor([0.93937271], requires_grad=True)
>>> print(qml.draw(qnode)(0.3))
 0: ──RX(0.35)──╭C─────────┤ ⟨Z⟩
 1: ────────────╰X──╭C─────┤
 2: ────────────────╰X──H──┤

A more flexible way to apply this is as a function, where we must specify the argument like so (supposing that circuit is defined as above, but without the decorator):

>>> transformed_circuit = overrotate_rx(0.1)(circuit)
>>> qnode = qml.QNode(transformed_circuit, dev)
>>> print(qml.draw(qnode)(0.3))
 0: ──RX(0.4)──╭C─────────┤ ⟨Z⟩
 1: ───────────╰X──╭C─────┤
 2: ───────────────╰X──H──┤

Quantum function transforms can also be used within other quantum functions. For example, define the following ansatz:

def ansatz(angle):
    qml.RX(angle, wires=0)
    qml.CNOT(wires=[0, 1])
    qml.CNOT(wires=[1, 2])
    qml.CNOT(wires=[2, 0])

Let’s use it within the context of a larger circuit, but with the overrotate_rx transform applied.

def nested_circuit(x, overrot_angle):
    qml.Hadamard(wires=0)
    qml.RX(x[0], wires=1)
    overrotate_rx(overrot_angle)(ansatz)(x[1])
    return qml.expval(qml.PauliZ(0))

Note the syntax of how the transform is applied here. First, overrotate_rx(overrot_angle) yields a quantum function transform. This is then applied to ansatz, yielding a new quantum function. That new function then accepts the input parameters x[1]. (For more a more detailed explanation, see the “Usage Details” section of the documentation).

Let’s create and run a QNode to inspect the circuit:

>>> x = np.array([1.6, -0.8])
>>> nested_qnode = qml.QNode(nested_circuit, dev)
>>> nested_qnode(x, 0.05)
tensor(-0.02919952, requires_grad=True)
>>> print(qml.draw(nested_qnode)(x, 0.05))
 0: ──H────────RX(-0.75)──╭C──────╭X──┤ ⟨Z⟩
 1: ──RX(1.6)─────────────╰X──╭C──│───┤
 2: ──────────────────────────╰X──╰C──┤

The RX gate within the transformed ansatz has been over-rotated, but not the first RX gate.

Differentiating quantum function transforms

Earlier, we mentioned that transforms preserve the differentiability of quantum functions. Transforms can also be differentiated using different machine learning frameworks, such as Torch and TensorFlow. To remain differentiable, any internal classical processing in the transform must be written using your framework of choice. Alternatively, the qml.math module provides a dispatch mechanism to support all frameworks.

For example, let’s rewrite our transform that overrotates RX by adding the square root of an input parameter.

@qml.qfunc_transform
def overrotate_rx(tape, overrot_angle=0):
    for op in tape.operations + tape.measurements:
        if op.name == "RX":
            qml.RX(op.parameters[0] + qml.math.sqrt(overrot_angle), wires=op.wires)
        else:
            qml.apply(op)

Next we construct a simple circuit that performs a single RX:

def simple_rotation(x):
    qml.RX(x, wires=0)
    return qml.expval(qml.PauliZ(0))

We choose an arbitrary value for the rotation and inspect the ideal circuit:

>>> x = 1.6
>>> dev = qml.device('default.qubit', wires=1)
>>> ideal_qnode = qml.QNode(simple_rotation, dev)
>>> print(qml.draw(ideal_qnode)(x))
0: ──RX(1.6)──┤ ⟨Z⟩

Now we can construct a noisy version of this circuit that is over-rotated by parameter 0.05:

>>> noisy_qfunc = overrotate_rx(0.05)(simple_rotation)
>>> noisy_qnode = qml.QNode(noisy_qfunc, dev)
>>> print(qml.draw(noisy_qnode)(x))
0: ──RX(1.82)──┤ ⟨Z⟩

The full value of the noisy angle is:

>>> x + np.sqrt(0.05)
1.823606797749979

By virtue of the differentiability of transforms, we can construct a new transform that learns how to undo this overrotation, by inserting an additional RX operation after each existing one:

@qml.qfunc_transform
def undo_rx_overrotation(tape, undo_angle=0):
    for op in tape.operations + tape.measurements:
        if op.name == "RX":
            qml.apply(op)
            qml.RX(undo_angle, wires=op.wires)
        else:
            qml.apply(op)

By default, the angle correction has a value of 0:

>>> new_qfunc = undo_rx_overrotation()(noisy_qfunc)
>>> new_qnode = qml.QNode(new_qfunc, dev)
>>> print(qml.draw(new_qnode)(x))
0: ──RX(1.82)──RX(0)──┤ ⟨Z⟩

Let’s learn what it should be. First, we need to define a cost function that compares the fixed, over-rotated result with the true result. We’ll choose a simple least-squares cost.

def cost(angle):
    new_qfunc = undo_rx_overrotation(angle)(noisy_qfunc)
    new_qnode = qml.QNode(new_qfunc, dev)
    return (ideal_qnode(x) - new_qnode(x)) ** 2

Now we can construct an optimizer and learn the best angle of the undo_rx_overrotation transform for this case:

>>> param = 0
>>> opt = qml.GradientDescentOptimizer()
>>> for step in range(200):
...     param = opt.step(cost, param)

The resulting angle is:

>>> print(param)
-0.223067860068242

We see that this is correct by adding it to the original angle—we should recover something close to our ideal value of x.

>>> x + np.sqrt(0.05) + param
1.600538937681737

Once learned, this correction transform can be used for arbitrary circuits in which the RX get over-rotated by the same amount.

Conclusion

These examples are just a few of the countless ways that transforms can be used in PennyLane. The qml.transforms module also contains examples of QNode transforms, as well as special tape transforms that not only return tapes, but also classical processing functions that act on the results obtained from executing those tapes. We hope the new decorators make it easier for you to write (and contribute) your own transforms—we look forward to seeing what you come up with!