PennyLane Torch Implementation Question

[johnmarktaylor91]

Hello, I had a somewhat specific question. I maintain a library called TorchLens that allows users to visualize the computational graph of PyTorch models and extract any intermediate activations. It seems to be having a hard time with PennyLane layers so I'm trying to understand why. For example, if I visualize the computational graph of a model there appears to be a gap between the input and output, suggesting that some tensor operations aren't being tracked in the middle.

My question is--does PennyLane's PyTorch implementation only use base PyTorch operations under the hood for tensor-processing operations, or does it define new functions for performing functions on tensors? I ask because TorchLens works by temporarily decorating all functions in the PyTorch namespace such that their execution is tracked and logged, and if PennyLane performs operations on tensors that don't boil down to base PyTorch functions (e.g., that are implemented with new C code) that would explain the issue. Many thanks for any insight on this issue.

qnn.pdf

[josh146]

Hi @johnmarktaylor91! Can I ask what PennyLane device are you using, and what differentiation method you are using?

The answer to the above depends on the underlying differentiation method --- if you are using a simulator that is written in PyTorch (such as default.qubit) and the differentiation method is set to backprop, then it should be end-to-end functions on Torch tensors.

However, if you are using a non-PyTorch compatible device, or different differentiation method (such as parameter-shift or adjoint), then PennyLane is defining a custom VJP rule for the quantum device/simulator, and the quantum execution will be a black box from Torch's perspective (it could be using NumPy, C++, or even executed on real quantum hardware).

[johnmarktaylor91]

Got it, thanks so much for the fast response! I am totally new to PennyLane so I really appreciate the help. I am running it on my normal Linux computer. Here's the code I ran--based on this, it sounds like it should be all PyTorch functions under the hood, right? If so, it sounds like that is not the source of the issue and I'll keep digging ;)

import torchlens as tl
import torch
import pennylane as qml
from os.path import join as opj

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

@qml.qnode(dev, diff_method="backprop")
def qnode(inputs, weights):
    qml.AngleEmbedding(inputs, wires=range(n_qubits))
    qml.BasicEntanglerLayers(weights, wires=range(n_qubits))
    return [qml.expval(qml.PauliZ(wires=i)) for i 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(2, 2)
softmax = torch.nn.Softmax(dim=1)
layers = [clayer_1, qlayer, clayer_2, softmax]
model = torch.nn.Sequential(*layers)

model_history = tl.log_forward_pass(model, torch.rand(1, 2), vis_opt='unrolled',
                                    vis_outpath=opj('/home/jtaylor/projects/torchlens/local_jmt/Debugging Scripts/qnn'))

[johnmarktaylor91]

In case it's helpful, here is the visualization of the computational graph I'm getting (green node is input, red node is output, other nodes are tensor operations, boxes correspond to modules)--it's still showing a hole in the middle. But know that only pytorch operations are involved will help me a lot with debugging.

qnn.pdf

[CatalinaAlbornoz]

Hi @johnmarktaylor91 ! Thank you for sharing your code. Let me investigate and get back to you in the next couple of days.

[johnmarktaylor91]

Thanks so much :)

[CatalinaAlbornoz]

Hi @johnmarktaylor91 , I'm having some trouble getting your code to run. I get an error RecursionError: maximum recursion depth exceeded while calling a Python object. I'll give it a try in a new environment and let you know if I can figure it out.

[johnmarktaylor91]

Thanks so much--please don't sink too much time into this, only if it's a quick fix. Much obliged :)

[CatalinaAlbornoz]

Hi @johnmarktaylor91 , I've got some good and bad news. The good news is that in a new environment I'm not getting the error that I was getting before so I can run your code with no problems. The bad news is that I don't get any output at all.

I haven't dug deep into the Torchlens documentation but do you have any pointers on what I may be doing wrong? I get no error at all but no pdf either. I'm on a mac.

[johnmarktaylor91]

Did you set vis_outpath? This sets the filepath for the visualization. Thanks for all your careful work on this, it is hugely appreciated!

[CatalinaAlbornoz]

Hi @johnmarktaylor91 , I did but I had a different issue with my environment. I've now managed to fix the environment's issue so now I do get the output! I'll test it out with different PennyLane configurations and let you know what I find.

[johnmarktaylor91]

Thanks so much :)

[CatalinaAlbornoz]

Hi @johnmarktaylor91,

I've tried several PennyLane configurations and I still don't know where the black box situation occurs.

I tried even the simplest circuit and the gap that you noticed is still there. If you run the circuit below (where nothing at all is happening inside) the problem is still there.

I'll check with the team if someone has capacity to take a look at this with fresh eyes.

qml.logging.enable_logging()

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

@qml.qnode(dev, diff_method="backprop")
def qnode(inputs, weights):
    #print(inputs)
    #qml.RX(inputs[0][0],wires=0)
    #qml.RY(weights,wires=0)
    return [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))] 

weight_shapes = {"weights": 1}
qlayer = qml.qnn.TorchLayer(qnode, weight_shapes)

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

model_history = tl.log_forward_pass(model, data, vis_opt='unrolled',
                                    vis_outpath=...)

[johnmarktaylor91]

Thanks so much--if it ends up being impractical please don't move heaven and earth from this, I really appreciate all your help so far!

[CatalinaAlbornoz]

I think Torchlens is very cool and I'd love to see it working with PennyLane. However the answer to the mystery is escaping me at the moment. I'm hoping to get some volunteers to take a look at it in about 2 weeks. I'll keep you updated on any findings!

[johnmarktaylor91]

Thanks so much 🙏 Really appreciate all your help!

[CatalinaAlbornoz]

No problem @johnmarktaylor91 !