Add support for measuring multiple observables

[jogisuda]

I am trying to set up my custom decoder, where I want to return the expectation of Z in multiple qubits, but one at a time (like this tutorial from PennyLane). The problem is that the outputs have no gradients, so I can't perform any kind of backward computation for training. The observables are like:

Hamiltonians = [Z0 @ I1 @ I2 @ ... @ In,
I0 @ Z1 @ I2 @... @ In,
I0 @ I1 @ I2 @ ... @ Zn]

For this I have the following snippet, where first I have an auxiliary function to build the list of observables:

import numpy as np


from qiboml.models.encoding import PhaseEncoding
from qibo import Circuit, gates
from qiboml.models.decoding import QuantumDecoding
from qibo.symbols import Z
from qibo.hamiltonians import SymbolicHamiltonian
from qibo import set_backend

set_backend("qiboml",platform="pytorch")

def create_hamiltonians(n_qubits):

    measurements = []
    for i in range(n_qubits):
        hamiltonians = [I(i) for i in range(n_qubits)]
        hamiltonians[i] = Z(i)
    
        H = 1
        for term in hamiltonians:
            H *= term
    
        H = SymbolicHamiltonian(H)
        measurements.append(H)
    
    return measurements

and then I also have a function to create a quantum net (encoding and circuit are returned):

def create_quantum_net(n_qubits, nlayers):
    
    qubits = list(range(n_qubits))
    
    # define the encoding
    encoding = PhaseEncoding(nqubits=n_qubits)

    q_weights = params.reshape(nlayers, n_qubits, 2)
    # build the computation circuit
    circuit = Circuit(n_qubits)

    print("\n[*] Creating a circuit with {} layers..\n".format(nlayers))
    for layer in range(nlayers):
        for q in qubits:
            circuit.add(gates.RY(q, theta=0.3, trainable=True))
            circuit.add(gates.RZ(q, theta=0.4, trainable=True))

        if n_qubits > 1:
            for i, q in enumerate(qubits[:-2]):
                circuit.add(gates.CNOT(q0=q, q1=qubits[i + 1]))
            circuit.add(gates.CNOT(q0=qubits[-1], q1=qubits[0]))

    print("Creating circuit with {} qubit...".format(n_qubits))
    return encoding, circuit

Finally, a class for creating a custom decoder:

class MyCustomDecoder(QuantumDecoding):

    def __init__(self, nqubits: int):
        super().__init__(nqubits)
        # build the observables using qibo's SymbolicHamiltonian
        self.nqubits = nqubits
        self.hamilts = create_hamiltonians(nqubits)

    def __call__(self, x: Circuit):
        # execute the circuit and collect the final state
        state = super().__call__(x).state()

        # calculate the expectation values
        return torch.tensor([self.hamilts[i].expectation(state) for i in range(self.nqubits)])
        
    # specify the shape of the output
    @property
    def output_shape(self) -> tuple[int]:
        return (self.nqubits)

And to define and run the model I have:

n_qubits = 2
q_depth = 2

decoding = MyCustomDecoder(nqubits=nqubits)

encoding, circuit = create_quantum_net(n_qubits, q_depth)

# encapsulate encoding, circuit and the multiple hamiltonians in QiBo syntax:
quantum_model = QuantumModel([encoding, circuit], decoding)

model = torch.nn.Sequential(
    quantum_model,
)
outputs = model(torch.randn(2))
print(outputs)

And output gives me:

tensor([-0.8379,  0.8457], dtype=torch.float64)

Whereas a simple model in PyTorch with one linear layer would give as output:

tensor([ 0.0198, -0.0311], grad_fn=<ViewBackward0>)

[BrunoLiegiBastonLiegi]

Hey, thanks for trying this out. I tested your code and for me it works fine until you do the explicit call to torch.tensor(...).
I think you should just replace that with torch.vstack or similars:

class MyCustomDecoder(QuantumDecoding):

    def __init__(self, nqubits: int):
        super().__init__(nqubits)
        # build the observables using qibo's SymbolicHamiltonian
        self.nqubits = nqubits
        self.hamilts = create_hamiltonians(nqubits)

    def __call__(self, x: Circuit):
        # execute the circuit and collect the final state
        state = super().__call__(x).state()

        # calculate the expectation values
        # torch.tensor apparently discards the gradient, you could use torch.vstack or 
        # other concatenations depending on the shape you need
        # return torch.tensor([self.hamilts[i].expectation(state) for i in range(self.nqubits)] 
        return torch.vstack([self.hamilts[i].expectation(state) for i in range(self.nqubits)])
        
    # specify the shape of the output
    @property
    def output_shape(self) -> tuple[int]:
        return (self.nqubits)

[jogisuda]

Thank you! It looks like creating a tensor makes torch detach the old tensors and hence lose the gradient information. stacking was the appropriate choice.