Changing a circuit meaurement once a Qnode is compiled

Hello! I have a use case where during the training of a hybrid quantum/classical model at the end of an epoch I would like to change the measurement that occurs at the end of a quantum circuit. Here is a basic hybrid model to exemplify.

import pennylane as qml
import torch

n_qubits = 2
dev = qml.device("default.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.sample()

n_layers = 2
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)

During the training of this model (with some loss function/data) at the end of the epoch I would instead like to measure the probability of each computational basis using qml.probs(). Ideally, the function would behave as the following I believe:

  • Accepts the Qnode as an argument

  • Replaces the measurement at the end to qml.probs()

  • Executes the circuit to return the [len(2**n_qubits] probability list

  • Replaces this measurement with the original qml.sample() so at the beginning of the next epoch, training continues as previous

Any help would be greatly appreciated, whether using a method like this or something simpler. It seems that once you’ve compiled the Qnode, accessing the operations and measurements externally is difficult to do.

Hey @Aaron_Thomas! You can return multiple things from a QNode like so:

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

@qml.qnode(dev)
def qnode():
    return qml.state(), qml.probs()

qnode()
(tensor([1.+0.j, 0.+0.j, 0.+0.j, 0.+0.j], requires_grad=True),
 tensor([1., 0., 0., 0.], requires_grad=True))

That said, you’ll run into issues when you sandwich this into a hybrid model. Another option would be this:

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

def return_a_qnode(measurement, **kwargs):

    @qml.qnode(dev)
    def qnode(inputs, weights):
        qml.AngleEmbedding(inputs, wires=range(n_qubits))
        qml.BasicEntanglerLayers(weights, wires=range(n_qubits))
        return measurement(**kwargs)

    return qnode

n_layers = 2
weight_shapes = {"weights": (n_layers, n_qubits)}

qlayer = qml.qnn.TorchLayer(return_a_qnode(qml.probs, wires=[0]), weight_shapes)
print(qlayer(torch.tensor([1.0, 1.0])))
tensor([0.4500, 0.5500], grad_fn=<ViewBackward0>)
clayer_1 = torch.nn.Linear(2, 2)
clayer_2 = torch.nn.Linear(2, 2)
softmax = torch.nn.Softmax()
layers = [clayer_1, qlayer, clayer_2, softmax]
model = torch.nn.Sequential(*layers)

model(torch.tensor([1.0, 1.0]))
tensor([0.6436, 0.3564], grad_fn=<SoftmaxBackward0>)

You can change the measurement by simply giving a different value to the measurement argument (which must be a measurement process: Measurements — PennyLane 0.34.0 documentation). That’s not 100% what you want but it might get you somewhere!

Other than that, would you be able to accomplish what you want to do with a transform? qml.transforms — PennyLane 0.34.0 documentation

Hopefully one of these options sparks something for you :grin:

Hi @isaacdevlugt,
Thank you for your response, your suggestion of a transformation was what I was thinking. I believe I’ve written the correct transformation but not sure in what way to apply it, if you could provide any insight that’d be incredible.

I have defined a class for my network, something like the following, I use torchlayer as later I connect with a classical network, though that’s not necessary for this example:

import pennylane as qml
from pennylane.qnn import TorchLayer 
from torch import Tensor

class Network(nn.Module):
    def __init__(
        self,
    ) -> None:
        super(Binary_Generator, self).__init__()

        self.n_qubits = 4
        self.depth = 2
        self.device = qml.device('default.qubit' , wires = self.n_qubits)

        q_weight_shapes = {"q_weights_y": (self.depth, self.n_qubits)}

        self.qnode = qml.QNode(self._circuit, device = self.device,
                                      interface="torch")

        self.q_network = TorchLayer(self.qnode, q_weight_shapes)

    def _circuit(self, inputs, q_weights_y, q_weights_z):
        """Builds the circuit to be fed to the connector as a QML node"""

        # Embedding layer
        AngleEmbedding(inputs, wires=range(self.n_qubits), rotation="X")

        for i in range(self.depth):
            for y in range(self.n_qubits):
                qml.RY(q_weights_y[i][y], wires = y)
            for y in range(self.n_qubits - 1):
                qml.CNOT(wires=[y, y + 1])

        return qml.sample()

    def forward(self, inputs: Tensor):
        return self.q_generator(inputs)

I then create an instance of this class for training later on

q_network = Network()

Like i mentioned i want to later change what is being measured during training so the function ive written is the following, which change the number of shots and the measurment from qml.sample() -> qml.probs():

from typing import Callable, Sequence, Union
import copy
from pennylane.tape import QuantumTape, QuantumScript
from pennylane.transforms.core import transform

@transform
def get_probslist(tape: QuantumTape) -> tuple[Sequence[QuantumTape], Callable]:

    tape = copy.deepcopy(tape)

    measurement = qml.probs()
    shots = 10000

    new_tape = QuantumScript(tape.operations, measurements = measurement, 
                                shots=shots, trainable_params = tape.trainable_params)
    
    def processing_fn(res):
        return res[0]

    return [new_tape], processing_fn

I want to apply the function to an instance of the class during function. This, i believe, has to happen at the Qnode level so my idea was to do the following:

#Some point during training

q_node = q_network.qnode
output = get_probslist(q_node)

However this doesn’t quite work, for one I can see the object is stored at the same memory point as the original qnode even though i deepcopy the model. The function is correct but im not sure how to apply once ive got an instance of the class. Any help would be greatly appreciated!

I’ll have to get back to you! Need to consult some folks internally. Sit tight!

1 Like

Hi @Aaron_Thomas! Sorry for the late reply.

Here I think a transform isn’t the right way to go. Think of a transform as a recipe that takes in one circuit and gives you (a) multiple circuits, (b) a recipe to combine execution results from those circuits. You can get a few more details here.

I had a quick attempt at getting this to work and it’s not so easy when there is a QNode wrapped up in a TorchLayer, as is the case here. One option might be to have a global SAMPLING boolean and to return qml.sample() if SAMPLING else qml.probs() in your QNode. You can then switch the boolean depending on what stage you are in the training process.

Hi @Tom_Bromley, thanks for the reply. This idea makes sense however whilst changing to qml. probs, I’d also need to change the shot count at the same time. Now I can use your logic to do both but the shots are baked into the qml.device() object. If I change both, specifically the shots as well, will that essentially restart the qml.device and clear any parameters?

Thanks @Aaron_Thomas! In PennyLane, shots can be set either when instantiating a device or when evaluating a QNode:

dev = qml.device("default.qubit", shots=100)

@qml.qnode(dev)
def f(x):
    qml.RX(x, 0)
    return qml.expval(qml.PauliZ(0))
>>> f(0.3)  # Executes with 100 shots
0.94
>>> f(0.3, shots=10)  # Overrides the device shots value to do 10 shots instead
1.0

Unfortunately, the TorchLayer wrapper is a bit less flexible, losing the ability to override the device shots value at QNode runtime. It’s something we can fix, but would need to fit it on the roadmap. For now, if you can avoid the TorchLayer wrapper for this usecase that’d be great, otherwise I can only recommend slightly hacky options like setting the (non-public) _shots attribute on the device:

>>> dev._shots = 1000  #!! Proceed with caution and don't rely on this functionality
>>> f(0.3)  # Now runs with 1000 shots
0.962

Hi @Tom_Bromley thank you for your response! This is what i was looking for however for my use case i need to use the torchlayer. If there was the ability in the future to change at runtime the shots when a qnode is wrapped in a torchlayer thatd be amazing! Thank you for the help

2 Likes