Pennylane + Braket with Adjoint method not working

PennyLane Help
1

Hello! I am trying to run a QML model with Pennylane + Keras on Braket. I am running into an error where it cannot use the adjoint method however, it complains about my measurements (Pauli Z across all wires). I got the code for the Braket device from this guide, but my measurement is a little different from the guide.

The following code (using random data) produces the error for me:

from braket.devices import Devices
import pennylane as qml
import tensorflow as tf
from tensorflow_addons.metrics import RSquare
from keras.layers import Dense
from tensorflow.keras.metrics import MeanSquaredError, MeanAbsoluteError
import numpy as np

example_X_train = np.array([[-1.86997793e+01,  4.02454270e+00,  2.76800223e+00,
         -2.68606211e+01,  1.37322594e+01, -3.68356304e+00],
        [ 1.87651854e+01,  9.24642075e+01,  5.25006713e+00,
         -1.93293771e+00, -1.18398824e+01,  1.25684404e+01],
        [ 7.80031328e-01,  4.91078049e+00,  4.66351891e-03,
          1.60988466e-02, -4.78166848e-01,  1.21960608e+00],
        [-1.21131161e+01, -9.68673860e+00, -3.43639144e+00,
         -5.80915398e+01,  4.26303421e+00, -3.96819741e+00],
        [ 1.11943184e+02, -4.23489171e+01, -3.86753766e+01,
          3.21862875e+00,  1.93067517e+01,  1.93340072e+01]])
example_Y_train = np.array([[ 981.46497203],
        [ 193.1293178 ],
        [1162.23848003],
        [ 500.02628397],
        [1327.104     ]])

n_qubits = 6
learning_rate=0.1
epochs=10
batch_size = 5

diff_method = 'device'
dev = qml.device(
    "braket.aws.qubit",
    device_arn="arn:aws:braket:::device/quantum-simulator/amazon/sv1",
    wires=n_qubits,
    shots=0,
)

@qml.qnode(dev, diff_method=diff_method)
def circuit(inputs, weights):
    qml.templates.AngleEmbedding(inputs, wires=range(n_qubits), rotation='X', id='X')
    # Ansatz
    [qml.RX(phi=weights[n], wires=n, id=weights[n]) for n in range(n_qubits)]
    return [qml.expval(qml.PauliZ(i)) for i in range(n_qubits)]

weight_shapes = {"weights": (n_qubits)}

qlayer = qml.qnn.KerasLayer(circuit, weight_shapes, output_dim=n_qubits, dtype='double')
clayer_out = tf.keras.layers.Dense(1, activation = "linear", dtype='double')
model = tf.keras.models.Sequential()

model.add(qlayer)
model.add(clayer_out)
model.build(input_shape=(None, 6))
model.summary()

opt = tf.keras.optimizers.Adam(learning_rate=learning_rate)
model.compile(opt, loss='mse',metrics=[RSquare(),MeanSquaredError(),MeanAbsoluteError()])

hist = model.fit(example_X_train, example_Y_train, epochs=epochs, verbose=1, batch_size=batch_size)

I get the error:

{..}.local/lib/python3.10/site-packages/braket/pennylane_plugin/braket_device.py:858: UserWarning: This circuit cannot be differentiated with the adjoint method. Falling back to the parameter-shift method, which will execute 12 circuits. To use the adjoint gradient method, make sure the circuit's only measurement is an expectation of one observable. To avoid gradient calculation, mark parameters with `requires_grad=False`.
  warnings.warn(

I have Pennylane version 0.34.0 and amazon-braket-pennylane-plugin version 1.24.4.

Thank you for the help!

2

Hey @vhazoglou!

The issue is here:

make sure the circuit's only measurement is an expectation of one observable.

You’re calculating the expval of more than one observable:

    return [qml.expval(qml.PauliZ(i)) for i in range(n_qubits)]

Let me know if that helps!

3

I eventually found a good solution to this.

Changing the measurement to one qubit wasn’t so much an option as it no longer will work well for this hybrid machine learning model.

Instead I just needed to switch to using SPSA here:

@qml.qnode(dev, diff_method='device')

should instead be

@qml.qnode(dev, diff_method='spsa')

And voilà, it works. :slight_smile:

This is also better for QML as it will cause fewer circuit executions.

One can read more about parameter-shift and SPSA in Pennylane here:

1 Like