How to create your own device in PennyLane

Maurice Weber (Xanadu resident)

PennyLane is designed to be hardware and device agnostic, allowing users to easily dispatch quantum circuits on different quantum devices. In addition to using built-in simulators, external backends can be accessed using plugins. PennyLane even allows you to create your own custom devices which is what we will cover in this how-to.

Devices are among the core objects in PennyLane and are primarily used to execute quantum circuits. You can use ideal state-vector simulators such as default.qubit, density matrix simulators for simulating noisy quantum circuits like default.mixed, and real quantum hardware.

Devices work in tandem with QNode objects and enable you to execute an arbitrary quantum circuit defined as a simple Python function. In this how-to, we will explain the ins and outs of setting up your own device, and how it can be used to execute an arbitrary quantum circuit.

The base class for all devices is the Device class which provides standardized methods for any type of system, regardless of whether it uses the qubit, qutrit, or continuous variable model for quantum computation. The base class for every qubit-based device is the QubitDevice class which is a child of the Device class and inherits its methods and properties. In addition to standardized methods implemented in its parent, the QubitDevice class also has a dedicated execute method and provides implementations for computing measurement statistics like marginal probabilities, expectation values, and variances.

Note: In the following example, we consider a qubit-based device and define a device class that inherits from the ‘QubitDevice’ class. If you want to define continuous-variable-based devices you should use ‘Device’ as the parent class.

We start by importing our favourite libraries:

import pennylane as qml
from pennylane import numpy as np

Now, let’s create our first minimal device, which is an ideal state-vector simulator with a single qubit, performing ideal computations. While, for the sake of simplicity, we fix the number of wires to one in this tutorial, this is generally not required and Device classes can be created to accept a variable to set the number of wires. The first step is to create a subclass of QubitDevice (which we will name MyDevice) and provide it with the required attributes:

class MyDevice(qml.QubitDevice):
    name = 'Ideal Single Qubit State-Vector Simulator'
    short_name = 'custom.qubit'
    pennylane_requires = '>=0.23'
    version = '0.0.1'
    author = 'Maurice'

    operations = {"RX", "RY", "RZ", "PhaseShift"}
    observables = {"PauliX", "PauliZ"}

Whenever we create a device, we have to provide it with these attributes. So let us go through them one by one:

  • name is the name of the device. You are free to choose whatever name you like here. We used ‘Ideal Single Qubit State-Vector Simulator’ in our example device as it is a concise description of the device.
  • short_name is used to instantiate the corresponding device class. After the device has been installed, users can then create a device instance via qml.device('custom.qubit', wires=num_wires).
  • author is the author of the device. This is Maurice in our example device.
  • pennylane_requires sets the PennyLane version required for the device. Note that you can use standard pip style version ranges.
  • version is the version of the custom device. In our case, we just set it to 0.0.1 since it is the first device of its kind.

Finally, the operations and observables attributes are defined as Python sets and are of particular importance to the device as they define the sets of operations and observables which are supported by the device.

Next, we provide the __init__ method which sets the number of wires to 1 and the shots to None since we simulate quantum computations analytically. Since we are creating our custom device, we can assume, for example, that our qubit instead of starting at \(|0 \rangle\), will start at \(|1\rangle\) by default! So, in the definition of our device, we set the initial state to be \(|1\rangle\).

def __init__(self):
    super().__init__(wires=1, shots=None)

    # create the initial state
    self._state = np.array([0, 1]) # |1❭ := (0, 1)

    # create a variable for future copies of the state
    self._pre_rotated_state = None

Each device class in PennyLane has a method called capabilities which returns a dictionary of key-value pairs indicating what the device can (and cannot) do. For example, the capabilities method of the QubitDevice returns the following dictionary:

{'model': 'qubit',
'supports_broadcasting': False,
'supports_finite_shots': True,
'supports_tensor_observables': True,
'returns_probs': True}

Since this corresponds to the parent class of our custom device, we need to overwrite the method in order to account for the capabilities of our own device. In particular, our device does not support finite shots or tensor products of observables (as it is a single qubit device). To do this, we update the capabilities as follows:

@classmethod
def capabilities(cls):
    capabilities = super().capabilities().copy()
    capabilities.update(
        returns_state=True,
        supports_finite_shots=False,
        supports_tensor_observables=False
    )
    return capabilities

where we have also recorded the capability of our device being able to return the current state. This is necessary for the method access_state in the QubitDevice parent, which allows returning a reduced density matrix of the state. In addition, for the access_state method to work, the device also needs to be able to return the state, which we implement as a class property returning the state prior to the rotations coming from previous measurements:

@property
def state(self):
    return self._pre_rotated_state

Furthermore, for devices based on the QubitDevice parent, we also need to implement the apply method. Since we are simulating an ideal device, we apply operations by simply multiplying the state with the matrix representation of the operations:

def apply(self, operations, rotations=None, **kwargs):
    for op in operations:
        # We update the state by applying the matrix representation of the gate
        self._state = qml.matrix(op) @ self._state

For observables for which it is required to make measurements on a basis different from the computational basis, e.g. PauliX, the state needs to be rotated before the measurement. For this reason, we also supply the apply method with optional argument rotations which rotate the state into the required basis. In our analytic device, we implement this simply via an additional set of matrix multiplications:

# store the pre-rotated state
self._pre_rotated_state = self._state.copy()

# apply the circuit rotations
for rot in rotations or []:
    self._state = qml.matrix(rot) @ self._state

Note that the implementation of the execute method in the QubitDevice class calls the apply method internally and automatically takes care of the rotations needed to measure a specific observable.

To make the last piece of our custom device, we need to implement the analytic_probability method since we are constructing an ideal device and the method is not implemented in the parent (see the PennyLane documentation for the abstract placeholder method in the parent). This method is expected to return the (marginal) probabilities of the computational basis states.

To that end, we first split the state into its real and imaginary parts and then use the marginal_prob method to get the marginal probabilities over the wires.

Note that applying the marginal_prob function here is actually unnecessary since we only have a single qubit device. In general, however, one might want to get the probabilities for a particular subset of qubits. Since we’re building a single qubit state-vector device, we get the analytic probabilities for the computational basis states by summing up the squares of the real and imaginary parts (elementwise) in the state-vector:

def analytic_probability(self, wires=None):
    if self._state is None:
        return None

    real = self._real(self._state)
    imag = self._imag(self._state)
    prob = self.marginal_prob(real**2 + imag**2, wires)
    return prob

Internally, if for example qml.expval is requested, the device will use these probabilities, together with the eigenvalues of the observable, to compute the final expectation value.

Finally, we provide a method that resets the device to its initial state, using the reset method

def reset(self):
    """Reset the device"""
    self._state = np.array([0, 1])

Putting the pieces together, we now have our final custom single qubit device:

class MyDevice(qml.QubitDevice):
    name = 'Ideal Single Qubit State Vector Simulator'
    short_name = "custom.qubit"
    pennylane_requires = ">=0.23"
    version = "0.0.1"
    author = "Maurice"

    operations = {"RX", "RY", "RZ", "PhaseShift"}
    observables = {"PauliX", "PauliZ"}

    def __init__(self):
        super().__init__(wires=1, shots=None)

        # create the initial state
        self._state = np.array([0, 1])

        # create a variable for future copies of the state
        self._pre_rotated_state = None

    @property
    def state(self):
        return self._pre_rotated_state

    @classmethod
    def capabilities(cls):
        capabilities = super().capabilities().copy()
        capabilities.update(
            returns_state=True,
            supports_finite_shots=False,
            supports_tensor_observables=False
        )
        return capabilities

    def apply(self, operations, rotations=None, **kwargs):
        for op in operations:
            # We update the state by applying the matrix representation of the gate
            self._state = qml.matrix(op) @ self._state

        # store the pre-rotated state
        self._pre_rotated_state = self._state.copy()

        # apply the circuit rotations
        for rot in rotations or []:
            self._state = qml.matrix(rot) @ self._state

    def analytic_probability(self, wires=None):
        if self._state is None:
            return None

        real = self._real(self._state)
        imag = self._imag(self._state)
        prob = self.marginal_prob(real ** 2 + imag ** 2, wires)
        return prob

    def reset(self):
        """Reset the device"""
        self._state = np.array([0, 1])

The common devices available in PennyLane are installed and registered as entry points in the setup.py file, which allows users to initialize devices via their short_name. However, as a handy feature, it is also possible to directly use your custom device by instantiating the class. That is, instead of defining it using its short name (e.g., qml.device("custom.qubit")), we use the device class directly:

device = MyDevice()

Then, we can define a quantum function in the usual way with the qml.qnode decorator

@qml.qnode(device=device)
def quantum_function():
    qml.Hadamard(wires=[0])
    return qml.expval(qml.PauliX(wires=[0]))

and call it like a usual Python function to evaluate it:

quantum_function()
>> -1.0

Remember that we are starting at \(\vert 1 \rangle\) and we have just calculated \(\langle 1 \vert HXH \vert 1 \rangle\). Doing the math we can see that it is correct since:

$$\langle 1 \vert HXH \vert 1 \rangle = \langle 1 \vert Z \vert 1 \rangle = -\langle 1 \vert 1 \rangle = -1.$$

And that is already all! Now we know how to define our own devices in PennyLane and use them to evaluate circuits. Also, creating your own device class is the first step if you want to build a full-fledged plugin to provide e.g., additional custom observables or operations, or to allow an external quantum library to take advantage of PennyLane’s automatic differentiation capabilities. You can visit the PennyLane documentation for more in-depth explanations about how to create your own plugin in PennyLane and/or install your custom devices.

Maurice Weber

Maurice Weber

Maurice is a PhD student at the Department of Computer Science at ETH Zurich. His research is focused on robustness guarantees for classical and quantum machine learning. Currently, he is a summer resident at Xanadu.