Example1_VQC_training.py

# Copyright 2025 Xanadu Quantum Technologies Inc.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# author: Joseph Lee (Xanadu Quantum Technologies Inc.)
"""
PennyLane example demonstrating VJP (vector-Jacobian product) + JAX
"""
import sys
import pennylane as qml
import jax.numpy as jnp
from timeit import default_timer as timer

# For pennylane==0.43 , this requires jax==0.6.2
import jax

# Set device based on command line argument
device_name = "lightning.amdgpu"

if len(sys.argv) > 1:
    if sys.argv[1].lower() == "cpu":
        device_name = "lightning.qubit"
    elif sys.argv[1].lower() == "gpu" or sys.argv[1].lower() == "amdgpu":
        device_name = "lightning.amdgpu"
    else:
        print(f"Unknown device {sys.argv[1]}, using GPU (lightning.amdgpu)")


# enable 64 bit precision for JAX
jax.config.update("jax_enable_x64", True)

# Try and scale up these numbers!
wires = 20
layers = 2

# Set a seed
key = jax.random.PRNGKey(0)

# Set number of runs for timing averaging
num_runs = 3

# Use device_name == `lightning.qubit` for CPU
# Use device_name == `lightning.amdgpu` for AMD GPU
dev = qml.device(device_name, wires=wires)


# Create QNode of device and circuit
@qml.qnode(dev)
def circuit(parameters):
    qml.StronglyEntanglingLayers(weights=parameters, wires=range(wires))
    return qml.expval(qml.PauliZ(wires=0))


# Set trainable parameters
shape = qml.StronglyEntanglingLayers.shape(n_layers=layers, n_wires=wires)
weights = jax.random.uniform(key, shape)

# Generate a random cotangent vector (matching output shape of circuit)
key, subkey = jax.random.split(key)
cotangent_vector = jax.random.uniform(subkey, shape=())


# Define VJP wrapper
def compute_vjp(params, vec):
    primals, vjp_fn = jax.vjp(circuit, params)
    return vjp_fn(vec)[0]


# JIT the VJP function
jit_vjp = jax.jit(compute_vjp)

# Warm-up run
_ = jit_vjp(weights, cotangent_vector).block_until_ready()

timing = []
vjp_res = None

for t in range(num_runs):
    start = timer()
    # Pass weights and the random cotangent vector
    vjp_res = jit_vjp(weights, cotangent_vector).block_until_ready()
    end = timer()
    timing.append(end - start)

print("Circuit measurements: \n", circuit(weights), "\n")

## Optional: print out circuit diagrams for small circuits
if wires < 5 and layers < 3:
    print("Raw circuit:\n", qml.draw(circuit)(weights), "\n")
    print("Expanded circuit:\n", qml.draw(circuit, level="device")(weights), "\n")

print("VJP Result (Gradient scaled by random vector): \n", vjp_res[0], "\n")
print("Mean timing: ", qml.numpy.mean(timing), "\n")