Obtaining circuit gradient with Torch

[vincentmr]

This is a follow-up on this question and this answer.

This time, let's take Torch, and the same simple circuit. It applies two gates to one qubit and evaluates the expectation value of PauliZ. When optimizing with Torch, the expectation value runs aways to -Inf. I think my issue is I cannot correctly select or filter which gates are parametrized (according to Torch). I have tried going around apply_to_arrays to pick and choose what gets requires_grad=True and update the torch.nn.ParameterDict field based on that. But then the optimizers do not work.

This may be an issue with my usage of Torch, but I was wondering if someone knew a quick way to fix this. Feel free to ignore if this is outside the scope of this forum.

# %%
import quimb as qu
import quimb.tensor as qtn
import numpy as np
import torch


# %%
def convert(x, requires_grad=False):
    return torch.tensor(x, dtype=torch.complex128, requires_grad=requires_grad)


quimb_obs = convert(qu.pauli("Z"), requires_grad=False)
where = 0
N = 1


# %%
def circuit(params):
    circ = qtn.Circuit(N)

    # hamadard on one of the qubits
    circ.apply_gate("RX", params[0], 0, parametrize=True)
    circ.apply_gate("RY", params[1], 0, parametrize=True)
    return circ


def loss_fn(ket):
    bra = ket.H
    ket = ket.gate(quimb_obs, 0)
    expec = bra & ket
    return np.real(expec.contract(all))


params = np.array([0.011, 0.012])
circ = circuit(params)
psi = circ.psi
loss_fn(psi)

# %%
for t in psi[0]:
    t.set_params(convert(t.get_params(), isinstance(t, qu.tensor.PTensor)))
# psi.apply_to_arrays(convert)


# %%
class TNModel(torch.nn.Module):
    def __init__(self, tn):
        super().__init__()
        # extract the raw arrays and a skeleton of the TN
        params, self.skeleton = qtn.pack(tn)
        # n.b. you might want to do extra processing here to e.g. store each
        # parameter as a reshaped matrix (from left_inds -> right_inds), for
        # some optimizers, and for some torch parametrizations
        self.torch_params = torch.nn.ParameterDict(
            {
                # torch requires strings as keys
                str(i): torch.nn.Parameter(initial, requires_grad=initial.requires_grad)
                for i, initial in params.items()
            }
        )

    def forward(self):
        # convert back to original int key format
        params = {int(i): p for i, p in self.torch_params.items()}
        # reconstruct the TN with the new parameters
        psi = qtn.unpack(params, self.skeleton)
        # isometrize and then return the energy
        return loss_fn(psi)


# %%
model = TNModel(psi)
model()

#%%

optimizer = torch.optim.Adagrad(model.parameters())
its = 100

for _ in range(its):
    optimizer.zero_grad()
    loss = model()
    loss.backward()
    optimizer.step()

[jcmgray]

Maybe doing something like the following would work:

constant_tn, variable_tn = tn.partition(['RX', 'RY'])

and then in you only pack/unpack the variable part, and call

psi = constant_tn | variable_tn

in your loss function?

Adding tag filtering to the pack/unpack functions would be nice to do at some point, this is the other major part of what TNOptimizer handles - i.e. selecting and grouping which tensors are actually being optimized.

(by the way I edited your Q just to add the ```python syntax highlighting, hope that's okay.)

[vincentmr]

Thanks for the suggestion. I have modified the code around the loss_fn as follows

params = np.array([0.011, 0.012])
circ = circuit(params)
psi = circ.psi
constant_tn, variable_tn = psi.partition(["RX", "RY"])
constant_tn.apply_to_arrays(lambda x: convert(x, requires_grad=False))
variable_tn.apply_to_arrays(lambda x: convert(x, requires_grad=True))


# %%
def loss_fn(tn):
    ket = copy.deepcopy(tn)
    ket &= constant_tn
    bra = ket.H
    ket = ket.gate(quimb_obs, where)
    expec = bra & ket
    return expec.contract(backend="torch").real

# [...]

# %%
model = TNModel(variable_tn)
model()

It seems to get me further, but I still have an error

    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

I think I'll move on for now, since I have a way forward with JAX (thanks to you). I'll revisit this in a bit and post the solution when I have found it.

(by the way I edited your Q just to add the ```python syntax highlighting, hope that's okay.)

Of course!

[jcmgray]

I haven't time to run through your example properly, but:

1. If you are using the Model infrastructure, I don't think you need to set requires_grad yourself, that is all handled by setting attributes as parameters.
2. You could use requires_grad, then simply contract the tensor network, and call x.backward() yourself, no need for torch.nn.Module or any abstractions at all:

import torch
import quimb.tensor as qtn

tn = qtn.TN_rand_reg(3, 2, 2)

for i in range(3):
    if i == 1:
        # or with get_params/set_params to handle PTensors
        tn[f'I{i}'].modify(apply=lambda x: torch.tensor(x, requires_grad=True))
    else:
        tn[f'I{i}'].modify(apply=lambda x: torch.tensor(x))
        
x = tn ^ ...
x.backward()

tn['I1'].data.grad
# tensor([[-1.8986, -1.4572],
#         [-0.0318,  0.0149]], dtype=torch.float64)

[vincentmr]

Thanks, this is what I was looking for. I tried modifying (and simplifying) the original example accordingly.

for i in range(len(psi[0])):
    if isinstance(psi[0][i], PTensor):
        psi[0][i].set_params(torch.tensor(psi[0][i].get_params(), requires_grad=True))
    else:
        psi[0][i].modify(apply=lambda x: torch.tensor(x, requires_grad=False))
tn = psi.H & psi.gate(quimb_obs, where)
x = tn ^ ...
x.backward()

tn[0][1].data.grad
tn[0][1].get_params().grad

The statement psi[0][i].modify(apply=lambda x: torch.tensor(x, requires_grad=False)) breaks the code, giving me the same error as before

RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

while if I set requires_grad=True in the same statement, grad is None.

[jcmgray]

This seems to be a bug. I don't know if this was always the case, but at the moment calling torch.tensor([[c, -s], [s, c]]) where all of those are scalars with gradients, doesn't propagate requires_grad, and this is the main step of creating the parametrized gate.

[jcmgray]

As a side note, there are lots of nice ways to access the tensors that would be recommended over

for i in range(len(psi[0])):
    psi[0][i]

i.e. getting the ith tensor tagged with site tag 0 each time. You can just do:

psi.select(["RY", "RZ"], "any").apply_to_arrays(to_variable)
psi.select(["RY", "RZ"], "!any").apply_to_arrays(to_constant)

etc.!

[vincentmr]

This seems to be a bug.

Would you like me to report the issue?

As a side note, there are lots of nice ways to access the tensors that would be recommended over

Thanks for the advice. I'll try to get rid of these operator[] accesses in the future.