How to train the rotation parameters of LCU of two QSP operation sequences?

Hello! I am trying to code linear combination of unitaries (LCU) of two quantum signal processing (QSP) operation sequences. With the code below, which I took from Pennylane’s demo of function fitting using QSP and I modified a little, I think I have succeeded in coding the LCU of two QSP operation sequences.

import torch
torch.set_default_dtype(torch.float64)  # Set default to double
import pennylane as qml
import math
import matplotlib.pyplot as plt

def rotation_mat(a):
    """Given a fixed value 'a', compute the signal rotation matrix W(a).
    (requires -1 <= 'a' <= 1)
    diag = a
    off_diag = (1 - a**2) ** (1 / 2) * 1j
    W = [[diag, off_diag], [off_diag, diag]]

    return W

def generate_many_sro(a_vals):
    """Given a tensor of possible 'a' vals, return a tensor of W(a)"""
    w_array = []
    for a in a_vals:
        w = rotation_mat(a)

    return torch.tensor(w_array, dtype=torch.complex64, requires_grad=False)

def unitaries(phi, W, wire):
    for angle in phi[:-1]:
        qml.RZ(angle, wires=wire)
        qml.QubitUnitary(W, wires=wire)
    qml.RZ(phi[-1], wires=wire)

def QSP_circ(phi_1, phi_2, W):
    """This circuit applies the the LCU of U_1 and U_2. See Xanadu Codercise section H.6 for reference of the
    qml.Hadamard(wires=0)  # set initial state |+>
    qml.ControlledQubitUnitary(qml.matrix(unitaries)(phi_1, W, 1), control_wires=[0], wires=1, control_values=0)
    qml.ControlledQubitUnitary(qml.matrix(unitaries)(phi_2, W, 1), control_wires=[0], wires=1, control_values=1)
    qml.Hadamard(wires=0)  # change of basis |+> , |->

d = 5
a_vals = torch.linspace(-1, 1, 50)
w_mats = generate_many_sro(a_vals)

gen = torch.Generator()
gen.manual_seed(444422)  # set random seed for reproducibility

for i in range(10):
    phi = torch.rand(d + 1, generator=gen) * 2 * torch.tensor([math.pi], requires_grad=False)
    matrix_func = qml.matrix(QSP_circ)
    y_vals = [matrix_func(phi[:2], phi[2:], w)[0, 0].real for w in w_mats]
    plt.plot(a_vals, y_vals, label=f"poly #{i}")

plt.vlines(0.0, -1.0, 1.0, color="black")
plt.hlines(0.0, -1.0, 1.0, color="black")

To train the parameter angle of RZ, I used the code below but I got an error

torch_pi = torch.Tensor([math.pi])

class QSP_Func_Fit(torch.nn.Module):
    def __init__(self, degree_1, degree_2, num_vals, random_seed=None):
        """Given the degree and number of samples, this method randomly
        initializes the parameter vector (randomness can be set by random_seed)
        if random_seed is None:
            self.phi = torch_pi * torch.rand(degree_1 + degree_2 + 2, requires_grad=True)
            gen = torch.Generator()
            self.phi = torch_pi * torch.rand(degree_1 + degree_2 + 2, requires_grad=True, generator=gen)
        self.phi = torch.nn.Parameter(self.phi)
        self.num_phi_1 = degree_1 +1
        self.num_vals = num_vals

    def forward(self, omega_mats):
        y_pred = []
        generate_qsp_mat = qml.matrix(QSP_circ)

        for w in omega_mats:
            u_qsp = generate_qsp_mat(self.phi[:self.num_phi_1], self.phi[self.num_phi_1:], w)
            P_a = u_qsp[0,0]
        return torch.stack(y_pred, 0)

class Model_Runner:
    def __init__(self, model, degree_1, degree_2, num_samples, x_vals, process_x_vals, y_true):
        """Given a model and a series of model specific arguments, store everything in
        internal attributes.
        self.model = model
        self.degree_1 = degree_1
        self.degree_2 = degree_2
        self.num_samples = num_samples

        self.x_vals = x_vals
        self.inp = process_x_vals(x_vals)
        self.y_true = y_true

    def execute(self, random_seed=191199, max_shots=2000, verbose=True):
        """Run the optimization protocol on the model using Mean Square Error as a loss
        function and using stochastic gradient descent as the optimizer.
        model = self.model(degree_1=self.degree_1, degree_2=self.degree_2, num_vals=self.num_samples, random_seed=random_seed)
        model = model.double()

        criterion = torch.nn.MSELoss(reduction="sum")
        optimizer = torch.optim.SGD(model.parameters(), lr=1e-5)

        t = 0
        loss_val = 1.0

        while (t <= max_shots) and (loss_val > 0.5):
            self.y_pred = model(self.inp)
            if t == 1:
                self.init_y_pred = self.y_pred
            loss = criterion(self.y_pred, self.y_true).double()
            loss_val = loss.item()

            if (t % 1000 == 0) and verbose:
                print(f"iter: {t}, loss: {round(loss_val, 4)} loss : {type(loss.item())}")

            t += 1
        self.model_params = model.phi 
        return [self.y_pred.tolist(), self.model_params.tolist()]

    def plot_result(self, show=True):
        plt.plot(self.x_vals, self.y_true.tolist(), "--b", label="target func")
        plt.plot(self.x_vals, self.y_pred.tolist(), ".g", label="optim params")
        plt.plot(self.x_vals, self.init_y_pred.tolist(), ".r", label="init params")

        if show:

import numpy as np

d1 = 4
d2 = 3  
num_samples = 50

def custom_poly_noisy(x):
    return torch.tensor(x**2 +  0.5*np.random.normal(0, 0.5, num_samples), requires_grad=False, dtype=torch.float)

a_vals = np.linspace(-1, 1, num_samples)
y_noisy = custom_poly_noisy(a_vals)

qsp_model_runner = Model_Runner(QSP_Func_Fit, d1, d2, num_samples, a_vals, generate_many_sro, y_noisy)

pred = qsp_model_runner.execute()

The error is

RuntimeError: Found dtype Float but expected Double

referring the line with loss.backward(). I have checked the original code and the loss value is indeed a float, but there is no error. I think the data type might not be the problem. Is there anyone who can help me?

Thank you very much. Below is the output of qml.about()

Name: PennyLane
Version: 0.32.0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
License: Apache License 2.0
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: PennyLane-Lightning

Platform info:           Linux-6.2.0-39-generic-x86_64-with-glibc2.37
Python version:          3.11.4
Numpy version:           1.23.5
Scipy version:           1.11.3
Installed devices:
- default.gaussian (PennyLane-0.32.0)
- default.mixed (PennyLane-0.32.0)
- default.qubit (PennyLane-0.32.0)
- default.qubit.autograd (PennyLane-0.32.0)
- default.qubit.jax (PennyLane-0.32.0)
- (PennyLane-0.32.0)
- default.qubit.torch (PennyLane-0.32.0)
- default.qutrit (PennyLane-0.32.0)
- null.qubit (PennyLane-0.32.0)
- lightning.qubit (PennyLane-Lightning-0.32.0)

Hey @verstrikt,

I think the problem could be this line’s dtype:

If you change that to a double, does that fix things?

I’m gonna try it. I will let you know if it works. Thank you!

It works! Thank you!!!

1 Like