How does vectorized_value_and_grad work?

[vincentmr]

I would to compute the gradient of certain circuits the way it's done in PennyLane. For instance, creating a circuit of strongly entangled layers and computing the expectation value of a local gate

# %% Import PennyLane
import pennylane as qml
import pennylane.numpy as np
import numpy as nnp

# %% Circuit parameters
n_wires = 3
n_layers = 4
obs = qml.PauliZ
where = 0
draw = False

# %% Create PennyLane circuit
dev = qml.device("lightning.qubit", wires=n_wires)


@qml.qnode(dev)
def circuit(parameters):
    qml.StronglyEntanglingLayers(weights=parameters, wires=range(n_wires))
    return qml.expval(obs(where))


shape = qml.StronglyEntanglingLayers.shape(n_layers=n_layers, n_wires=n_wires)
weights = np.random.random(size=shape)
if draw:
    print(qml.draw(circuit, expansion_strategy="device")(weights))

# %% Evaluate circuit and gradient
nparams = nnp.random.random(size=shape)
params = np.array(nparams)
if draw:
    print(qml.draw(circuit, expansion_strategy="device")(params))
expv_qml = circuit(params)
print("expv_qml = ", expv_qml)
grad_qml = qml.grad(circuit)(params)
print("grad_qml = ", grad_qml)

Analogously, with Quimb

# %% Import Quimb + defs
import quimb as qu
import quimb.tensor as qtn
import jax.numpy as jnp

quimb_obs = qu.pauli("Z")
where = 0
backend = "jax"


# %% Define Quimb circuit
def single_qubit_layer(circ, params, gate_round):
    """Apply a parametrizable layer of single qubit ``U3`` gates."""
    for i in range(circ.N):
        # initialize with random parameters
        circ.apply_gate(
            "RZ", params[i, 0], i, parametrize=True, contract=False, gate_round=gate_round
        )
        circ.apply_gate(
            "RY", params[i, 1], i, parametrize=True, contract=False, gate_round=gate_round
        )
        circ.apply_gate(
            "RZ", params[i, 2], i, parametrize=True, contract=False, gate_round=gate_round
        )


def two_qubit_layer(circ, gate_round):
    """Apply a layer of constant entangling gates."""
    offset = (gate_round % (n_wires - 1)) + 1
    for i in range(circ.N):
        circ.apply_gate("CNOT", i, (i + offset) % n_wires, contract=False, gate_round=gate_round)


def quimb_circuit(params):
    n_layers, n_wires, _ = params.shape
    circ = qtn.Circuit(N=n_wires)
    for j in range(n_layers):
        single_qubit_layer(circ, params[j, :, :], gate_round=j)
        two_qubit_layer(circ, j)
    return circ


jparams = jnp.array(nparams, dtype=jnp.float64)
circ = quimb_circuit(jparams)
if draw:
    circ.uni.draw(color=["RZ", "RY", "CNOT"])


# %% Run Quimb circuit
expv_qub = circ.local_expectation(quimb_obs, (where), backend=backend)
print(expv_qub)
print(f"Error: {np.abs(expv_qub-expv_qml):+1.3e}")


# %% Optim Quimb circuit


def fun_expect(psi):
    ket = psi
    bra = ket.H
    expec = bra & ket.gate(quimb_obs, where)
    return np.real(expec.contract(all, optimize="auto-hq", backend=backend))


expv_qub = fun_expect(circ.psi)
print(expv_qub)
print(f"Error: {np.abs(expv_qub-expv_qml):+1.3e}")

Up to this point, everything is in order. I would then like to obtain the gradient. This can be done with the parameter shift method as follows

# %% Param-shift

grad_qub = np.zeros_like(grad_qml, dtype=complex).ravel()
for i in range(jparams.size):
    tmp = np.array(jparams).ravel()
    tmp[i] -= np.pi / 2.0
    # circ0 = circuit(tmp.reshape(jparams.shape))
    circ0 = fun_expect(quimb_circuit(tmp.reshape(jparams.shape)).psi)
    tmp = np.array(jparams).ravel()
    tmp[i] += np.pi / 2.0
    # circ1 = circuit(tmp.reshape(jparams.shape))
    circ1 = fun_expect(quimb_circuit(tmp.reshape(jparams.shape)).psi)
    grad_qub[i] = (circ1 - circ0) / 2.0

grad_qub = grad_qub.reshape(jparams.shape)
print(f"Error: {np.max(np.abs(grad_qub-grad_qml)):+1.3e}")

which works but it inefficient. I am thus wondering how the get the gradient using a differentiable backend like JAX. I have tried using vectorized_value_and_grad from TNOptimizer

# %% Grad Quimb circuit


def loss(psi):
    return jnp.abs(fun_expect(psi))


tnopt = qtn.TNOptimizer(
    circ.psi,
    loss,
    tags=["RY", "RZ"],  # only optimize U3 tensors
    autodiff_backend="jax",  # use 'autograd' for non-compiled optimization
    optimizer="L-BFGS-B",  # the optimization algorithm
)
expv_qub = tnopt.vectorized_value(jparams.ravel())
print(expv_qml)
expv_qub, grad_qub = tnopt.vectorized_value_and_grad(jparams.ravel())
grad_qub = grad_qub.reshape(grad_qml.shape)
print(expv_qub, grad_qub)
print(f"Error: {np.max(np.abs(grad_qub-grad_qml)):+1.3e}")

, but I cannot get the correct value (unless all input parameters are the same) and gradient (in any case). I am thus wondering how to use vectorized_value_and_grad or what to use instead?

[jcmgray]

I haven't looked over your examples in detail, but my first thought would be that the vectorizer that quimb uses internally has no fixed relation to jparams.ravel(), other than the number of parameters, since it treats the tags and tensor network structure as the 'base' data-structure.

The process it follows is something like:

1. parse the TN, finding the tagged (/untagged) tensors, extracting either the raw data or parameters into a list, the order is determined from the graph and tag structure
2. then flatten and stack all of these into a vector, which is the vectorized parameter and gradient space

Generally the point of the TNOptimizer is to optimize the parameters inplace, so user access to the gradients is not supported. If you are just interested in the gradients then using TNOptimizer probably overcomplicates things!

The easiest thing to do is just use jax or torch normally, and treat quimb simply as a way to dispatch all the oprations which are then traced through. The function qtn.pack and qtn.unpack were written to help here, and full examples are here and here, though they also involve full optimizations not just getting the gradient.

If it was helpful, one might add a method to TNOptimizer that returns a new tensor network with the gradients rather than parameters injected, simply as a convenient way to access them in a labelled fashion.

[jcmgray]

Here's a small snippet that might help illustrate:

import jax
import quimb.tensor as qtn

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

# params is a dict of arrays
# skeleton is an 'empty' tensor network
params, skeleton = qtn.pack(tn)

@jax.jit
@jax.value_and_grad
def fn(params):
    # inject parameters into new tensor network
    tn = qtn.unpack(params, skeleton)
    # now we can just use quimb stuff as usual
    return tn ^ ...

fn(params)
# (Array(-6.866255, dtype=float32),
#  {0: Array([[[-1.7965045 , -4.473413  ],
#           [ 0.17244443, -7.8906946 ]],
#
#          [[-1.3230906 , -3.3164647 ],
#           [ 0.74837565, -7.5376062 ]]], dtype=float32),
#   1: Array([[[-0.9642699, -1.9810126],
#           [-1.7096559, -1.7649908]],
#
#          [[ 1.5658708,  3.390893 ],
#           [ 1.4973092,  1.5036218]]], dtype=float32),
#   2: Array([[[-8.8693905,  2.301557 ],
#           [ 2.5381231, -2.548389 ]],
# 
#          [[-3.3398378, -7.224867 ],
#           [-1.154387 ,  1.9598453]]], dtype=float32),
#   3: Array([[[-0.2630614 ,  0.41338658],
#           [-5.5481505 , 10.54403   ]],
#
#          [[-1.4778212 ,  0.13725984],
#           [-8.386109  , 11.030711  ]]], dtype=float32)})

you don't have to use pack/unpack, the point here is just that you define a 'raw' array function.

If fact for jax specifically you can use the even simpler (due to #148 + #151):

qtn.jax_register_pytree()

@jax.jit
@jax.value_and_grad
def fn(tn):
    return tn ^ ...

fn(tn)
# (Array(-6.866255, dtype=float32), TensorNetworkGen(tensors=4, indices=6))

and the gradients will be back in that returned TN.

[vincentmr]

Thanks for responding so quickly! As you pointed out, using TNOptimizer overcomplicates things. I was simply looking for a way to get the value and gradient of a circuit, so let's put it aside. The links you provided were useful. I am still having some issues extracting gradients however.

Let's take JAX, for example, and the following simple circuit. It applies two gates to one qubit and evaluates the expectation value of PauliZ. The cost function correctly return a value of 0.851540 but does not return the expected gradient [-0.51043865, -0.1026782 ] (it returns [ 1.70308117+0.j -0.23942441-1.02087731j]) instead. I am not really getting how and where to specify which parameters I'm interested in tracking, and how to fix this example.

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


# %%
quimb_obs = qu.pauli("Z")
where = 0
N = 1


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

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


# %%
import jax

qtn.jax_register_pytree()


@jax.jit
@jax.value_and_grad
def cost(psi):
    ket = psi
    bra = ket.H
    ket = ket.gate(quimb_obs, 0)
    expec = bra & ket
    return expec.contract(all, backend="jax")


params = np.array([0.54, 0.12])
value, grad = cost(circuit(params).psi)
grad = np.array(grad[0][0].data)
print(value, grad)

[jcmgray]

Yes so currently that function will take the gradient of the raw arrays of every tensor in psi. You need two tweaks:

1. Use circ.apply_gate("RX", params[0], 0, parametrize=True). This inserts PTensor instances instead of Tensor instances into the tensor network, these know that their data actually comes from f(params) and is not raw.
2. Only supply the actually varying tensors (i.e. the rotations) to the cost function.

It's a bit clunky currently (since quimb is not primarily a parametrized qu circ library, and haven't implemented this use-case before) but here is a working snippet:

import jax
import quimb as qu
import quimb.tensor as qtn
import numpy as np

quimb_obs = qu.pauli("Z")
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


params = np.array([0.54, 0.12])
circ = circuit(params)
psi = circ.psi


@jax.jit
@jax.value_and_grad
def cost(params):
    
    ket = psi.copy()
    ket["GATE_0"].set_params(params[0:1])
    ket["GATE_1"].set_params(params[1:2])
    
    bra = ket.H
    ket = ket.gate(quimb_obs, 0)
    expec = bra & ket
    return expec.contract(all, backend="jax").real


value, grad = cost(params)
value, grad
# (Array(0.85154057, dtype=float32),
#  Array([-0.5104387 , -0.10267819], dtype=float32))

A general interface for simply computing the gradients of specific tensors would probably be nice.

[vincentmr]

This works now. I'll go back and try to adapt this to my original example.

I'll mark you original answer as "the answer" since I think it correctly points out that using TNOptimizer is simply not the way to handle this. Thanks for the help here, much appreciated.

[jcmgray]

Nice, it needed a slight fix but you can also do this now:

import jax
import quimb as qu
import quimb.tensor as qtn
import numpy as np

qtn.jax_register_pytree()

quimb_obs = qu.pauli("Z")
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


params = np.array([0.54, 0.12])
circ = circuit(params)
psi = circ.psi.astype("complex64")


@jax.jit
@jax.value_and_grad
def cost(ket):
    bra = ket.H
    ket = ket.gate(quimb_obs, 0)
    expec = bra & ket
    return expec.contract().real


value, tn_grad = cost(psi)

tn_grad['RX'].get_params(), tn_grad['RY'].get_params()
# (Array([-0.5104387], dtype=float32), Array([-0.10267819], dtype=float32))

I.e. supply a TN with parametrized tensors and get the 'gradient TN' directly back. Note it this case all the tensors will have their gradients populated, not just the rotation tensors.

[vincentmr]

Looks really neat, thank you so much. To be clear,

Note it this case all the tensors will have their gradients populated, not just the rotation tensors.

do you mean any gate that has parametrize=True?

[jcmgray]

No I do mean all the tensors. parametrize=True means that the Circuit class puts a PTensor into the network, these tensors get differentiation w.r.t. parameters not the raw array, whereas the 'parameters' of a normal tensor are simply the raw array and these will get differentiated.

If you turn parametrize off in the above, you'll see that the gradient in tn_grad['RX'] is now the gradient of actual gate 2x2 matrix, so taking a step in that direction will not keep you in the unitary space!