AWS Quantum Technologies Blog

Explore quantum algorithms faster by running your local Python code as an Amazon Braket Hybrid Job with minimal code changes

Today, we are excited to announce the hybrid_job Python decorator, a new feature of the Amazon Braket SDK that enables quantum algorithm researchers to seamlessly execute local Python functions as an Amazon Braket Hybrid Job with only one additional line of code, making it easier to build, test, and run variational algorithms.

Researchers today are actively studying variational algorithms such as the quantum approximate optimization algorithm (QAOA), quantum machine learning algorithms, and the variational quantum eigensolver (VQE) to explore solutions to industry-relevant problems in several verticals, including finance, chemistry, pharmaceuticals, optimization, materials science.

To help accelerate this research, in 2021, we launched Amazon Braket Hybrid Jobs, a tool designed to make running quantum algorithms easier by offering customers priority access to quantum processors or QPUs to help minimize wait times in queues, along with the convenience of a “fire-and-forget” experience. Customers can simply submit their algorithm’s code to Amazon Braket, which manages the underlying compute to run iterative, variational algorithms. Braket releases the compute resources when the code terminates, so you only pay for what you use.

Until today, researchers would submit their Python code and dependencies in a container to the Hybrid Jobs API. However, most researchers do the bulk of their algorithm development in Jupyter notebooks. With this launch, we are now making the development, testing, and running of these algorithms even easier, by enabling customers to run Python functions directly as a hybrid job with only one additional line of code, helping them shorten the development lifecycle for building quantum algorithms.

As a bonus, it is now possible to use this feature with PennyLane, an open-source software framework for differentiable quantum programming designed to simplify the development of variational algorithms. PennyLane has a collection of tutorials (PennyLane demos), designed to introduce you to the basics of quantum programming, and to make accessible the latest theory, ideas, and algorithms in quantum computing. These tie explanations to code that you can copy, download, run, and extend, enabling you to turn ideas into experiments, and discover new ideas faster.

With this launch, we are adding the @hybrid_job Python decorator to PennyLane demos to make it even easier to go from running the demos on your local device to running them on managed simulators or QPUs available on Amazon Braket.

How it works

First, we import the necessary packages:

import pennylane as qml
from pennylane import numpy as np
from braket.jobs import hybrid_job
from braket.jobs.metrics import log_metric
from braket.devices import Devices

Next, we write our program as a Python function. To use this new capability, we add the @hybrid_job Python decorator. The Braket SDK automatically recognizes that this code should run as a Hybrid Job, and imports the code along with the necessary dependencies into a container.

Like all hybrid jobs, you’ll need to declare which device you would like the hybrid job to use, using the device argument in the @hybrid_jobdecorator. Quantum tasks that are created as part of the hybrid job get priority over other tasks waiting in the regular quantum task queue to enable your algorithm to run more efficiently and predictably without interruptions.

In this example, we set device=Devices.Rigetti.AspenM3 to get priority on our target QPU. For embedded simulators, when the simulator will run on the same classical host alongside the rest of your algorithm, you should set device=“local:<provider>/<simulator>”. For example, for the PennyLane Lightning CPU simulators, set device=”local:pennylane/lightning.qubit”.

device_arn = Devices.Rigetti.AspenM3

@hybrid_job(device=device_arn) # set priority QPU
def qpu_qubit_rotation_hybrid_job(num_steps=10, stepsize=0.5):
    # AWS devices must be declared within the decorated function.
    device = qml.device(
        "braket.aws.qubit",
        device_arn=device_arn.value,  # Make sure the device ARN matches the hybrid job device ARN
        wires=2,
        shots=1_000,
    )

    @qml.qnode(device)
    def circuit(params):
        qml.RX(params[0], wires=0)
        qml.RY(params[1], wires=0)
        return qml.expval(qml.PauliZ(0))

    opt = qml.GradientDescentOptimizer(stepsize=stepsize)
    params = np.array([0.5, 0.75])

    for i in range(num_steps):
        # update the circuit parameters
        params = opt.step(circuit, params)
        expval = circuit(params)

        log_metric(metric_name="expval", iteration_number=i, value=expval)

    return params

Now we create a hybrid job by calling the function as usual:

qpu_job = qpu_qubit_rotation_hybrid_job(num_steps=10, stepsize=0.5)
print(qpu_job)

This returns an AwsQuantumJob object that contains the device ARN, region, and job name, structured like this:

AwsQuantumJob(‘arn': 'arn: aws:braket: <aws-region>:<account_id>: job/qubit-rotation-hybrid-<timestamp>')

You can check whether a device is currently available, and get a sense for how long we will wait before the hybrid job runs using this snippet of code:

from braket.aws import AwsDevice

AwsDevice(device_arn).is_available
AwsDevice(device_arn).queue_depth().jobs

Once our job starts running, we can check the status of the job by reviewing the job.state(), or otherwise the Status column in the Braket console as shown in Figure 1.

Figure 1. Illustration showing the status of the running hybrid job in the Amazon Braket console.

Figure 1. Illustration showing the status of the running hybrid job in the Amazon Braket console.

For this example, it should take a few minutes to complete. When it’s done, we can see the results by running:

qpu_job.result()
An example of the output from a hybrid job, running in our Jupyter notebook.

An example of the output from a hybrid job, running in our Jupyter notebook.

Additionally, the Braket Hybrid Jobs automatically plots the metrics recorded during the execution of the program in the Braket console. The expectation value decreases with each iteration as you rotate the qubit.

Figure 2: Expectation value evaluated on a Rigetti Aspen M-3 for each iteration as shown from the Braket console.

Figure 2: Expectation value evaluated on a Rigetti Aspen M-3 for each iteration as shown from the Braket console.

In this example, we showed how you can take a Python function and run it as a hybrid job on the Braket service by adding only one line of code.

Running PennyLane Demos on Amazon Braket

Now, let’s see how someone starting from PennyLane Demos can use this functionality to go from running the demo locally on their laptop, to executing it on a managed simulator or QPU on Amazon Braket.

There are now over 110 demos on the PennyLane website, running the gamut from quantum hardware modalities, optimization, quantum chemistry, algorithms, quantum machine learning – and more. While it has always been easy to run the PennyLane Demos on Amazon Braket, the new @hybrid_job Python decorator is designed to make it even easier.

Starting today, on select PennyLane demos, a Run on Amazon Braket button will be available in the sidebar.

Figure 3. Example PennyLane demo showing the Run on Amazon Braket button in the Hardware section

Figure 3. Example PennyLane demo showing the Run on Amazon Braket button in the Hardware section

Clicking on this button will open a modal, with information about the Amazon Braket service. In addition, links to download the @hybrid_job version of the demo, are available. After clicking the link to view the Braket version of the demo, you’ll be directed to GitHub to view and download a notebook designed to work with Amazon Braket and the Hybrid Jobs feature.

Figure 4 Modal with information about the Braket service and links to helpful resources

Figure 4 Modal with information about the Braket service and links to helpful resources

From here, simply download and run the notebook from your local Jupyter environment or a Braket notebook instance. If you don’t already have Amazon Braket set up, refer to one of our previous blog posts to get started.

Conclusion

To accelerate research in quantum computing, customers need easier ways to get started running their variational quantum-classical workloads.

In this post we announced a new, simpler way for you to run hybrid algorithms using Braket Hybrid Jobs, by just adding one line of code to your local Python functions. This capability applies to any Python-based framework like the Braket SDK, PennyLane, or Qiskit. To illustrate this, we collaborated with PennyLane to expand a subset of PennyLane demos to include this functionality.

Check out our documentation, example notebooks, and the PennyLane demos to learn more.