Batch inputs and TorchLayer interactions with composite quantum circuit

Hello! I’m quite new to both PyTorch and PennyLane, but I’m trying to build a quantum autoencoder, following this article, for didactic purposes.

I’ve already written something but the more i debug the more it seems as everything fall apart. I reckoned one major problem of my code would be that i tried to implement the encoder, decoder and other parts of my code without using qml.qnn.TorchLayer, which only then I learned it’s been made for these kind of things. So I’m trying to implement the change, but I have structural doubts. You see, until now i structured my model as a sequence of five classes, namely Preprocessor, Encoder, Compressor, Decoder and Postprocessor, which expand PyTorch’s nn.Module (except the processors). Encoder, Compressor and Decoder are quantum, meaning that they’re made of a variational quantum circuit, with trainable parameters except for the input encoding.

As an example, i’ll paste here the code of the encoder (the others are analogous):

class Encoder(nn.Module):

    def __init__(self, nQubits, depth):
        super().__init__()

        self.nQubits = nQubits
        self.depth = depth
        self.rotParams = nn.Parameter(torch.rand(depth, nQubits, 3, device=tDevice), requires_grad=True)
        self.entParams = nn.Parameter(torch.rand(depth, nQubits, device=tDevice), requires_grad=True)

    def forward(self):
        self._circuit()

    def _circuit(self):
        for i in range(self.depth):
            self._repeatedLayer(layer=i)

    def _repeatedLayer(self, layer):
        # Rotations
        for i in range(self.nQubits):
            qml.Rot(self.rotParams[layer,i,0], self.rotParams[layer,i,1], self.rotParams[layer,i,2], wires=i)

        # Entanglement
        for i in range(self.nQubits-1):
            qml.CRZ(self.entParams[layer,i], wires=(i,i+1))
        qml.CRZ(self.entParams[layer,self.nQubits-1], wires=(self.nQubits-1,0))

The three quantum circuit would be then incorporated in an upper level class, which would run every circuit and return the results:

class MolQAE(nn.Module):

    def __init__(self, nQubits, latQubits, eDepth, dDepth):
        super().__init__()
        self.nQubits = nQubits
        self.latQubits = latQubits
        self.eDepth = eDepth
        self.dDepth = dDepth

        self.embedder = StateEmbedder(self.nQubits)
        self.encoder = Encoder(self.nQubits, self.eDepth)
        self.compressor = Compressor(self.nQubits, self.latQubits)
        self.decoder = Decoder(self.nQubits, self.dDepth)

        self.qnode = qml.QNode(self._circuit, qDevice, interface='torch', diff_method=diffMethod)
        self.qnode = TorchLayer(self.qnode)
        
    def forward(self, input):
        output = qml.snapshots(self.qnode)(input)
        initial = output['initialState']
        trash = output['trashState']
        final = output['execution_results']
        
        return initial, trash, final
    
    def _circuit(self, inputs):
        self.embedder(inputs)
        qml.Snapshot('initialState', measurement=qml.probs(wires=range(self.nQubits)))
        self.encoder()
        self.compressor()
        qml.Snapshot('trashState', measurement=qml.probs(wires=range(self.nQubits)))
        self.decoder()
        
        return qml.probs(wires=range(self.nQubits))

Here StateEmbedder is a class which handles the input state encoding.

So, this is what i’ve done, but as said everything seems to be wrong now that i need to add TorchLayer (i need it because i want proper batch handling and without it many inefficient for loops are to be used, as far as i understand).

For instances:

  • Am i supposed to not wrap into a qnode every circuit in encoder/compressor/decoder, but only the circuit in the upper level class? And if not, how i deal with the fact that i must return a measurement at the end of a node, but i need not?
  • At which level should i add TorchLayer? I mean, should i wrap Encoder, Compressor and Decoder as TorchLayers or it’s ok wrapping only MolQAE? Since my classes are separate, how should i deal with trainable parameters? I know TorchLayer automatically makes and stores its parameters (right?), but what if i have “nested” circuits?

I hope i wrote my doubts clear. Sorry if it’s a lot but it seems i just can’t get my head around it.

And thanks to anyone who will answer.

P.S.: Anyway, here’s qml.about():

Name: PennyLane
Version: 0.41.1
Summary: PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Train a quantum computer the same way as a neural network.
Home-page: https://github.com/PennyLaneAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /home/famiglia/miniconda3/lib/python3.13/site-packages
Requires: appdirs, autograd, autoray, cachetools, diastatic-malt, networkx, numpy, packaging, pennylane-lightning, requests, rustworkx, scipy, tomlkit, typing-extensions
Required-by: PennyLane_Lightning

Platform info:           Linux-6.14.0-29-generic-x86_64-with-glibc2.41
Python version:          3.13.2
Numpy version:           2.3.3
Scipy version:           1.16.0
Installed devices:
- lightning.qubit (PennyLane_Lightning-0.41.1)
- default.clifford (PennyLane-0.41.1)
- default.gaussian (PennyLane-0.41.1)
- default.mixed (PennyLane-0.41.1)
- default.qubit (PennyLane-0.41.1)
- default.qutrit (PennyLane-0.41.1)
- default.qutrit.mixed (PennyLane-0.41.1)
- default.tensor (PennyLane-0.41.1)
- null.qubit (PennyLane-0.41.1)
- reference.qubit (PennyLane-0.41.1)

P.P.S.: I’ll add my python notebook here, for a more comprehensive reference. You’ll find some comments in italian, my language, but don’t worry, it’s nothing important, they’re just annotations.

Notebook download

Hi @danim , welcome to the Forum!

Building a quantum autoencoder is not the easiest project, especially if you’re new to PyTorch and PennyLane. Since you mentioned you’re building this for didactic purposes I would actually recommend going through the PennyLane Codebook instead. It’s built with learners in mind! If you already know quantum computing (maybe a different SDK) I’d recommend starting with the PennyLane Fundamentals module. If you’re new to quantum computing, then you could start with the module on Introduction to Quantum Computing or the learning path on Fundamentals of Quantum Computing.

That being said, if you specifically want to write a quantum autoencoder I would recommend starting from the code in the ANNNI demo. It has a section on Quantum Anomaly Detection, which is the quantum version of an autoencoder.

Let us know if you have any further questions!

I hope this helps.