How to add custom gates and templates to PennyLane

Implementing your own gates in PennyLane does not require deep manipulation of the code base but is as easy as listing some core properties of the operation to be added. Even easier, you can construct your own templates to create structured special purpose circuits and simplify your workflow.

Adding a custom gate

We will build up an example of a custom gate, beginning with a minimal callable Python object and successively adding information to it.

The bare minimum

Let’s start by looking at a minimal callable implementation of a one-parameter two-qubit gate:

from pennylane.operation import Operation

class IncognitoGate(Operation):
    num_params = 1
    num_wires = 2
    par_domain = "R"

As we can see, from a programming perspective, we don’t need much to define a new custom gate in PennyLane (although it is not particularly useful at this stage). The minimal information required to form a valid, callable subclass of Operation is the following:

1. How many (variational) parameters does the gate take as input? (num_params)

2. On how many wires does the gate act? (num_wires)

3. From which domain are the parameters of the gate chosen? (par_domain)

   The available domains are:

• “N”: natural numbers and 0.
• “R”: real numbers, i.e. floating point numbers.
• “A”: arrays of real or complex values.
• “L”: lists of arrays of real or complex values.
• None: there are no parameters. Note that par_domain has to be provided explicitly in this case (as it is an abstract method).

A working example

While it forms a well-defined gate, the IncognitoGate does not do anything yet and in addition it will of course not be supported on any PennyLane device. Let us therefore promote it to a meaningful operation, for example a Pauli rotation generated by the Pauli word X\otimes X, which can be implemented via PauliRot:

import pennylane as qml
from pennylane.operation import Operation
from pennylane import numpy as np

class RXX(Operation):
    num_params = 1
    num_wires = 2
    par_domain = "R"

    @staticmethod
    def decomposition(theta, wires):
        return [qml.PauliRot(theta, 'XX', wires=wires)]

We add the method decomposition to relate our new gate to the gates that PennyLane already knows. This suffices to make the gate useful and usable in a quantum function:

dev = qml.device('default.qubit', wires=2, shots=100)

@qml.qnode(dev)
def circuit(theta):
    RXX(theta, wires=[0, 1])
    qml.Hadamard(1)
    return qml.expval(qml.PauliZ(0) @ qml.PauliY(1))
>>> circuit(0.3)
tensor(-0.12, requires_grad=True)

Notice how we did not have to implement the action of the gate in detail. The support for our brand-new gate RXX was simply inherited and all devices supporting qml.PauliRot - either directly or by decomposing it - will also support RXX.

Adding details

Now that we know the bare minimum required to create a useful gate, let’s consider some optional information we may want to provide about our gate:

Gradient details

As our new gate is parametrized by a real-valued parameter, we may want to tell PennyLane that there is a parameter-shift rule to compute its gradient analytically. This can be done by providing the subclass with the method grad_method. Available methods for differentiation are:

• “A”: Analytic differentiation via the parameter shift rule. The particular rule is stored in the grad_recipe of the gate and defaults to the two-term parameter shift rule.
• “F”: Numeric differentiation via finite difference. The order of the rule (1 for forward finite difference or 2 for central finite difference) and the shift size h can be passed to qml.qnode.
• None: No differentiation of the operation is possible/available.

If we set grad_method = "A" but the applicable parameter shift rule is not the default two-term rule, provide a grad_recipe to the gate explicitly. An example would be the four-term rule available as four_term_grad_recipe in PennyLane. Take a look at one of the gotchas regarding which gradient rule is used.

Matrix representation details

In addition to an implicit definition via decomposition it is often useful to provide the matrix of the gate explicitly. We do this via the class method _matrix which generates the class attribute matrix whenever we instantiate the new gate. Unless we register a new gate as supported by a specific device, the matrix will not be used internally, but it can be very useful in other contexts. For example, below we actually make use of the matrix attribute of another operator to define the generator property. (Also see a gotcha about only providing _matrix.)

Instead of _matrix, we may define the instance attribute matrix directly, using the property self.parameters (and self.wires if needed). For our rotation gate, this would take the following form:

@property
def matrix(self):
    c = np.cos(0.5 * self.parameters[0])
    s = np.sin(0.5 * self.parameters[0])
    return np.array(
        [
            [c, 0, 0, -s],
            [0, c, -s, 0],
            [0, -s, c, 0],
            [-s, 0, 0, c]
        ]
    )

Furthermore, for some optimizers the generating operator for the gate is used. For our gate, we provide PennyLane with this information via generator = [(qml.PauliX(0) @ qml.PauliX(1)).matrix, -0.5], where the first list entry is the operator and the second entry is the scaling prefactor. Alternatively, we could pass a NumPy array of the correct size as first entry, in which case we may include the scaling in the generator directly.

Other details

For self-adjoint operators or rotation gates like RXX we may want to tell PennyLane directly how to obtain the adjoint by implementing the adjoint class method.

The full gate

To sum it up, here is the full definition of our custom rotation gate:

import pennylane as qml
from pennylane.operation import Operation
from pennylane import numpy as np

class RXX(Operation):
    num_params = 1
    num_wires = 2
    par_domain = "R"

    grad_method = "A"
    grad_recipe = None # This is the default but we write it down explicitly here.

    generator = [(qml.PauliX(0) @ qml.PauliX(1)).matrix, -0.5]

    @staticmethod
    def decomposition(theta, wires):
        return [qml.PauliRot(theta, 'XX', wires=wires)]

    @staticmethod
    def _matrix(*params):
        theta = params[0]
        c = np.cos(0.5 * theta)
        s = np.sin(0.5 * theta)
        return np.array(
            [
                [c, 0, 0, -s],
                [0, c, -s, 0],
                [0, -s, c, 0],
                [-s, 0, 0, c]
            ]
        )

    def adjoint(self):
        return RXX(-self.data[0], wires=self.wires)

This new gate now can be used just like any other operation in the PennyLane stack, and will be supported on a wide range of devices.

Gotchas

Before we move on to creating custom templates in PennyLane, let us look at some gotchas when adding a custom gate:

1. If a gate has multiple parameters (i.e. num_params > 1), the parameter domain (par_domain) and the differentiation method (grad_method) of all parameters has to coincide.

2. The grad_recipe method has to be given per parameter. That means for a two-parameter gate that satisfies the two-term (four-term) rule with respect to the first (second) parameter, we would write

   grad_recipe = (
       [
           [0.5, 1, 0.5*np.pi],
           [-0.5, 1, -0.5*np.pi],
       ],
       qml.ops.qubit.four_term_grad_recipe[0],
   )

   As we can see, the structure is a tuple of single gradient recipes, which in turn are lists of entries. Each of these entries finally consists of the prefactor of the term, c_i, a prefactor for the gate parameter, a_i, and the parameter shift, s_i. These are used to compute

Regarding the stored four-term grad recipe, note that we had to extract the first (and only) entry of the tuple qml.ops.qubit.four_term_grad_recipe.

1. One could think about providing a custom gate with _matrix only, fully defining its action on a quantum state, and skip decomposition. However, this typically will fail because PennyLane devices do not know how to execute the gate yet. There are two ways to fix this:

   First, creating a new, special-purpose device would allow us to use the implementation via _matrix directly. Second, adding a custom gate to the native gate set of the default.qubit device simple and makes the gate available even without a decomposition:

   dev = qml.device('default.qubit', wires=3)
   dev.operations.add("RXX")

   Note that extending the native gate set in general is more complicated or might even be impossible (for example on quantum hardware devices).

2. If we do not register a custom gate with the device we are using (e.g. by creating a custom device), PennyLane will make use of the decomposition. When computing a derivative with respect to the gate parameter, the decomposition will be done first and therefore grad_method and grad_recipe of the custom gate will not be used at all.

3. For gates that do not use diff_method = "A", PennyLane will not allow the definition of a grad_recipe to make clear that it would not be used anyways.

4. In the future, the abstract method par_domain will be removed. This would not break our gates defined above but would make the definition of par_domain superfluous.

Adding a custom template

Let’s now look at how to build templates in PennyLane. Templates make our lives much easier by avoiding code duplication and simplifying the structure of our programs. For a guideline and technical details, visit the PennyLane documentation.

A minimal template

Let’s first look at a very simple template that prepares Bell states on neighbouring qubits.

Inheriting from Operation

Creating your own custom template in PennyLane and even adding it to the built-in template library can be done easily by implementing a qml.operation.Operation subclass, just like we did for the custom gates above. We will implement the action of the template within the class method expand. This method will be used internally by the PennyLane device to obtain a tape with known operations that then can be executed. This is the recommended way to construct templates which you plan on using often or even adding to the PennyLane library.

import pennylane as qml
from pennylane.operation import Operation, AnyWires

class BellPairLayer(Operation):
    num_params = 0 # The template does not take any parameters.
    num_wires = AnyWires # The template works for any number of wires.
    par_domain = None # We will not pass parameters to the template.

    def expand(self):
        # Create the tape with all operations.
        with qml.tape.QuantumTape() as tape:
            # Iterate over all qubit pairs
            for i in range(len(self.wires)//2):
                qml.Hadamard(wires=[self.wires[2*i]])
                qml.CNOT(wires=[self.wires[2*i], self.wires[2*i+1]])

        return tape

We can use the new template like any PennyLane Operation. For example, let’s prepare the Bell states and sum over all Pauli Z measurements.

from pennylane import numpy as np

num_wires = 8

# Instantiate a qubit device with num_wires=8 wires and 100 shots per expectation value.
dev = qml.device('default.qubit', wires=num_wires, shots=100)

# Simply add an expectation value measurement to the template and create a QNode.
@qml.qnode(dev)
def circuit():
    # As we want to call the template on all wires, we may use dev.wires.
    BellPairLayer(wires=dev.wires)
    return [qml.expval(qml.PauliZ(wire)) for wire in dev.wires]

# Sum the expectation values into a cost function.
def cost():
    return np.sum(circuit())
>>> cost()
tensor(-0.12, requires_grad=False)

We see that adding a template to PennyLane does not require much more than spelling out the used operations in the class method expand. The only additional information that is required are the abstract methods

• num_params, which often will be 0 as above or 1 because we unpack parameters within expand,
• num_wires, which usually will be set to the placeholder AnyWires to allow for any number of wires, and
• par_domain, which often will be "A" (or None for num_params=0), but does not play a major role and will be deprecated in future releases.

Using the template decorator

For simple cases like the Bell pair template above, there is an even easier way of defining the templates, namely as quantum function decorated with template:

import pennylane as qml

@qml.template
def bell_pair_layer(wires):
    # Iterate over all qubit pairs
    for i in range(len(wires)//2):
        qml.Hadamard(wires=[wires[2*i]])
        qml.CNOT(wires=[wires[2*i], wires[2*i+1]])

num_wires = 8
# Instantiate a default qubit device with num_wires=8 wires and 100 shots per expectation value
dev = qml.device('default.qubit', wires=num_wires, shots=100)
# Simply add an expectation value measurement to the template and create a QNode.
@qml.qnode(dev)
def circuit():
    # As we want to call the template on all wires, we may use dev.wires.
    bell_pair_layer(wires=dev.wires)
    return [qml.expval(qml.PauliZ(wire)) for wire in dev.wires]

# Sum the expectation values into a cost function.
def cost():
    return np.sum(circuit())

As we can see, the quantum function does not construct and return a tape like BellPairLayer.expand() did above, but calls the required operations directly. In addition, the template decorator allows us to apply built-ins like qml.inv() to the template.

Using the template afterwards is a piece of cake 🍰!

While the decorator method for defining a template is simpler for small circuits like the one shown here, it provides less capabilities and flexibility than the subclassing approach. This makes subclassing the preferred approach to provide templates to others and share it via the PennyLane template library. For more information on contributing your template, please refer to the corresponding technical guideline within the PennyLane documentation.