Qml.ctrl operation on GPU

Hello, I am trying to implement a quantum circuit that contains controlled CY gates. The circuit contains multiple layers. Following is the code snippet for a single layer.

def ansatz_2(n_qubits, y_weight):
    for wire, y_weight in enumerate(y_weight):
                qml.ctrl(op=qml.RY(y_weight, wires=wire), control=(wire-1)%n_qubits)

Other details are as follows:

dev = qml.device("default.qubit", wires=n_qubits)
interface = "torch"

When I run the ansatz on the CPU, it works well. However, on running the same ansatz on GPU, it throws the following error:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0, CPU!

Also, If I replace the layer with the following code snippet, the operation happens on GPU without any error.

def ansatz_2(n_qubits, y_weight):
    for wire, y_weight in enumerate(y_weight):
            qml.RY(y_weight, wires=wire)

So, I assume that I am doing something wrong with qml.ctrl implementation. Any help here would be appreciated.

Hey @imakash! I’m missing a lot of your code here, so I can’t say for certain what’s going on. That said, it sounds to me like you’re passing in two torch.tensor objects that are on different devices. Something akin to this:

import pennylane as qml
import torch

dev = qml.device("default.qubit", wires=2)

@qml.qnode(dev, interface="torch")
def circuit(weight_1, weight_2):
    qml.ctrl(op=qml.RY(weight_1, wires=1), control=0)
    qml.RX(weight_2, wires=0)
    return qml.state()

weights_1 = torch.tensor(0.1, device='cpu')
weights_2 = torch.tensor(0.1, device='cuda') # if I change this to 'cpu' it works

print(circuit(weights_1, weights_2))

Can you share your full code?

Hello @isaacdevlugt

Thank you for your response.

The entire code is quite complex and part of an ongoing research project. Therefore, I can’t post it on a public forum. However, are you able to run qml.ctrl operation on GPU?

Ah! Was able to replicate the issue. Had to dust off my other laptop with an NVIDIA card :sweat_smile:. You found a bug! Nice work! I encourage you to make a bug report on our GitHub (click here and it will direct you to the bug report form).

Indeed, this code replicates the error you see:

import pennylane as qml
import torch

dev = qml.device("default.qubit.torch", wires=2, torch_device='cuda')

@qml.qnode(dev, interface="torch")
def circuit(weight_1, weight_2):
    qml.ctrl(op=qml.RY(weight_1, wires=1), control=0) # when this gate is removed, all is fine
    qml.RX(weight_2, wires=0)
    return qml.state()

weights_1 = torch.tensor(0.1, device='cuda')
weights_2 = torch.tensor(0.1, device='cuda')

print(circuit(weights_1, weights_2))
# RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0, cpu!

As a fix for now, you can use qml.CRY (see here) instead of qml.ctrl. This works:

import pennylane as qml
import torch

dev = qml.device("default.qubit.torch", wires=2, torch_device='cuda')

@qml.qnode(dev, interface="torch", diff_method="backprop")
def circuit(weight_1, weight_2):
    qml.CRY(weight_1, wires=[0, 1])
    qml.RX(weight_2, wires=0)
    return qml.state()

weights_1 = torch.tensor(0.1, device='cuda')
weights_2 = torch.tensor(0.1, device='cuda')

print(circuit(weights_1, weights_2))
'''
 tensor([0.9988+0.0000j, 0.0000+0.0000j, 0.0000-0.0500j, 0.0000+0.0000j],
       device='cuda:0', dtype=torch.complex128)
'''
1 Like

Also worth mentioning here is that lightning.gpu is the best way to utilize a GPU in PennyLane. You can find out more here: Installation — PennyLane-Lightning-GPU 0.30.0.dev0 documentation

That said, if it doesn’t suit your needs, then no worries!

Our resident CUDA whisperer @mlxd has already made a fix for this: Qml.ctrl operation on GPU - #5 by isaacdevlugt

Hey @isaacdevlugt

Thank you very much for your help. I will also submit a bug report as you suggested.

1 Like

Awesome! Looks like this is set to be merged in good time!