Combining PennyLane with Amazon Braket

PennyLane integrates with Amazon Braket to add additional features for quantum machine learning and optimization. This introductory tutorial walks you through how to train a quantum circuit using Amazon Braket simulators and PennyLane's automatic differentiation capabilities.

Setup

PennyLane is already installed on Braket notebook instances. On a local machine, PennyLane can be installed by following these instructions. It can then be imported with:

To use Braket as a backend in PennyLane, we have to create a PennyLane device. Here we will first create a device that uses the local Braket simulator that runs on your local laptop (or on the server that hosts this notebook).

Below we will also show you how to scale out simulations to the AWS cloud.

Defining a circuit

We will choose a simple two-qubit circuit with two controllable rotations and a CNOT gate.

The qml.qnode(dev) decorator binds the circuit to the local Braket device. Now, every time that circuit() is called, the quantum computation defined in the function above will be executed with Braket.

Evaluating the circuit and accessing its gradient

Let's set some values for our controllable parameters:

The circuit can be evaluated with these parameters using

A crucial element of machine learning and optimization is accessing the gradient of a model with respect to its parameters. This functionality is built into PennyLane:

Here, dcircuit is a callable function that evaluates the gradient of the circuit, i.e., its partial derivatives with respect to the controllable parameters.

Training the circuit

Suppose we now want to minimize the output of the circuit by updating its parameters. This can be done using gradient-based optimization.

First, an optimizer is fixed:

The next step is to run the optimizer for a chosen number of iterations:

Running circuits on Braket's on-demand simulator, SV1

So far we have used the local Braket simulator. This is a great choice for quick prototyping, but it is not suitable for large circuits with many qubits and does not provide a connection to quantum hardware.

Amazon Braket also provides access to on-demand, high-performance simulators and quantum processing units (QPUs) from different providers. These devices can be accessed through PennyLane by changing a single line of code, unlocking the potential for machine learning and optimization on quantum hardware and high performance simulators!

Each remote Braket device can be selected through its ARN. The supported devices on Braket are listed here. For now, we will pick the on-demand SV1 simulator.

In PennyLane, all remote Braket devices are then accessed through a single PennyLane device named braket.aws.qubit.

A follow up tutorial shows you how to use the remote device to run multiple circuits in parallel, while the QAOA tutorial takes a deeper dive into graph optimization, including using SV1 to optimize a 20-node graph.

Let's execute our circuit on SV1, as well as calculating the gradient:

Running on a QPU with Amazon Braket Hybrid Jobs

In this notebook, the classical part of the algorithm was running locally. For longer-running algorithms or those requiring more intensive compute resources, it's recommended to dispatch the algorithm to Amazon Braket Hybrid Jobs, which fully manages the classical infrastructure, allowing you to focus on the algorithm. For example, you can train a larger circuit or increase the number of iterations.

The second benefit of running the algorithm as a hybrid job is that for iterative algorithms that require repeated calls to a QPU, you retain priority for that QPU. Once your quantum tasks are created in the hybrid job, they run ahead of other tasks waiting in the regular quantum task queue. This is because hybrid jobs have a separate queue from standalone tasks, ensuring that only a single hybrid job can run on a QPU at a time. This means your algorithm will not be interrupted by other quantum tasks, so it will run more efficiently and predictably. However, hybrid jobs have a separate queue from standalone tasks, so only a single hybrid job can run on a QPU at a time. For a single quantum circuit or a batch of circuits, it's recommended to create quantum tasks instead of hybrid jobs. Only iterative algorithms benefit from QPU priority queuing.

Note that hybrid jobs have at least a one-minute startup time since they create a containerized environment on Amazon EC2. So for very short workloads, there is likely no need to create a hybrid job.

You can run your local Python code as an Amazon Braket hybrid job by annotating your code with the `@hybrid_job`` decorator, as shown in the following code example. Only Python 3.12 is supported by default. For custom Python versions, you can choose to use a custom container from Amazon Elastic Container Registry (ECR) (see BYOC).

In the following code, we create a hybrid job for 10 iterations targeting Rigetti Cepheus-1-108Q. Since we specified Cepheus-1-108Q as the device, this job will run once Cepheus-1-108Q is available and has no jobs running ahead of it.

How long will it take the hybrid job to run? Let's first check if the device is currently available with AwsDevice(device_arn).is_available().

Next, we check the hybrid job queue depth with AwsDevice(device_arn).queue_depth().jobs.

If the device is available and there are no hybrid jobs currently running, then it should take about 5 minutes to complete.

Once the hybrid job has completed, we can retrieve the results with job.results().