Problems with gradient computing on pennylane lightning

Hey!
I am trying to train a hybrid torch model , it runs fine with default.qubit but with lightning.qubit or lightning.gpu or even qiskit.remote it fails , I observed two errors , and both of them originates from using qml.probs instead of qml.expval , the First error would be NotImplementedError: Computing the gradient of broadcasted tapes with respect to the broadcasted parameters using the parameter-shift rule gradient transform is currently not supported. See #4462 for details.
the code:

import pennylane as qml
import torch
import numpy as np
from sklearn.datasets import make_moons

X, y = make_moons(n_samples=200, noise=0.1)
y_ = torch.unsqueeze(torch.tensor(y), 1) # used for one-hot encoded labels
y_hot = torch.scatter(torch.zeros((200, 2)), 1, y_, 1)

n_qubits = 2
dev = qml.device(“lightning.qubit”, wires=n_qubits)

@qml.qnode(dev)
def qnode(inputs, weights):
qml.AngleEmbedding(inputs, wires=range(n_qubits))
qml.BasicEntanglerLayers(weights, wires=range(n_qubits))
return qml.probs(wires=[wire for wire in range(n_qubits)])

n_layers = 6
weight_shapes = {“weights”: (n_layers, n_qubits)}

qlayer = qml.qnn.TorchLayer(qnode, weight_shapes)

clayer_1 = torch.nn.Linear(2, 2)
clayer_2 = torch.nn.Linear(4, 2)
softmax = torch.nn.Softmax(dim=1)
layers = [clayer_1, qlayer, clayer_2, softmax]
model = torch.nn.Sequential(*layers)

opt = torch.optim.SGD(model.parameters(), lr=0.2)
loss = torch.nn.L1Loss()

X = torch.tensor(X, requires_grad=True).float()
y_hot = y_hot.float()

batch_size = 5
batches = 200 // batch_size

data_loader = torch.utils.data.DataLoader(
list(zip(X, y_hot)), batch_size=5, shuffle=True, drop_last=True
)

epochs = 6

for epoch in range(epochs):

running_loss = 0

for xs, ys in data_loader:
    opt.zero_grad()

    loss_evaluated = loss(model(xs) , ys)
    loss_evaluated.backward()

    opt.step()

    running_loss += loss_evaluated

avg_loss = running_loss / batches
print("Average loss over epoch {}: {:.4f}".format(epoch + 1, avg_loss))

y_pred = model(X)
predictions = torch.argmax(y_pred, axis=1).detach().numpy()

correct = [1 if p == p_true else 0 for p, p_true in zip(predictions, y)]
accuracy = sum(correct) / len(correct)
print(f"Accuracy: {accuracy * 100}%") 

the second error:

RuntimeError: Function ExecuteTapesBackward returned an invalid gradient at index 0 - got but expected shape compatible with [5]

it happens when i change the loss calculation method to be something custom-like, just as an exampleloss_evaluated = model(xs) * + entropy_loss where entopy_loss is a value computed from the output lets say (the error happens without it too) note that the [5] comes from the batch size , i read here a solution that i should write my own torch layer instead of using the existing one, but i am not really sure if that would fix this gradient computing problem. Thanks!

Platform info: Linux-6.10.11-100.fc39.x86_64-x86_64-with-glibc2.38
Python version: 3.12.4
Numpy version: 1.26.4
Scipy version: 1.13.1
Installed devices:

    lightning.qubit (PennyLane_Lightning-0.37.0)
    qiskit.aer (PennyLane-qiskit-0.37.0)
    qiskit.basicaer (PennyLane-qiskit-0.37.0)
    qiskit.basicsim (PennyLane-qiskit-0.37.0)
    qiskit.ibmq (PennyLane-qiskit-0.37.0)
    qiskit.ibmq.circuit_runner (PennyLane-qiskit-0.37.0)
    qiskit.ibmq.sampler (PennyLane-qiskit-0.37.0)
    qiskit.remote (PennyLane-qiskit-0.37.0)
    default.clifford (PennyLane-0.37.0)
    default.gaussian (PennyLane-0.37.0)
    default.mixed (PennyLane-0.37.0)
    default.qubit (PennyLane-0.37.0)
    default.qubit.autograd (PennyLane-0.37.0)
    default.qubit.jax (PennyLane-0.37.0)
    default.qubit.legacy (PennyLane-0.37.0)
    default.qubit.tf (PennyLane-0.37.0)
    default.qubit.torch (PennyLane-0.37.0)
    default.qutrit (PennyLane-0.37.0)
    default.qutrit.mixed (PennyLane-0.37.0)
    default.tensor (PennyLane-0.37.0)
    null.qubit (PennyLane-0.37.0)
    lightning.gpu (PennyLane_Lightning_GPU-0.37.0)
    cirq.mixedsimulator (PennyLane-Cirq-0.36.0.post0)
    cirq.pasqal (PennyLane-Cirq-0.36.0.post0)
    cirq.qsim (PennyLane-Cirq-0.36.0.post0)
    cirq.qsimh (PennyLane-Cirq-0.36.0.post0)
    cirq.simulator (PennyLane-Cirq-0.36.0.post0)
    qrack.simulator (pennylane-qrack-0.9.1)

Hi @Moustafa ,

I cannot guarantee that it will work but I strongly suggest that you follow the indication in my post from August 5th here. This should fix your first error. I would recommend starting with the code I shared and changing one element at a time instead of several at once. This can make it easier to identify and solve issues as they arise. Once you follow this approach of creating your own hybrid model let me know if you still get stuck on other issues.

Finally, I strongly suggest upgrading your PennyLane version to v0.38 (the latest one at the moment). I hope this helps!

Hey! @CatalinaAlbornoz
Yes, I have tried implementing the code that you mentioned, which indeed solved my first problem with lightning.qubit since the batchs are being processed one by one now, but a strange thing happened, the speed of my code when i use the custom layer is an order of magnitude slower than using TorchLayer with default.qubit, which made me go back to that setting.
any idea about why is this happening? , and if using TorchLayer in this context would give me good results even if it will use the default simulator and take a long time to execute? (TorchLayer isn’t broken, right? xd)

Thanks!

Hi @Moustafa,

TorchLayer isn’t broken as far as we know.

Different combinations of code + configuration + device will give different execution times, and sometimes they’re slow in all cases.

At the moment part of the work involved in developing simulations in quantum computing is finding the right configuration that will work for your specific problem. The selector at the bottom of our performance page can give you some insights on when to use each device, but ultimately it will depend on your specific use-case, and for you it might be that TorchLayer + default.qubit is the best option.