Second order differentiation using adjoint differentiation

Hi,

I noticed that calculating qml.jacobian(qml.grad(circuit)) using (lightning.qubit, diff_method=‘adjoint’) is not possible but using other diff methods is possible. But other diff methods are quite slow. Can someone please let me know if this is the case or am i making mistake.

Thanks

Hey @Abhishek! Welcome to the forum :smiley:!

Tagging @mlxd here because he’s the expert, but pennylane-lightning does support adjoint differentiation — there are caveats! For details, see here.

TLDR, adjoint differentiation has the following restrictions:

  • As it requires knowledge of the statevector, only statevector simulator devices can be used.
  • Only expectation values are supported as measurements.
  • Does not work for parametrized observables like Hamiltonian or Hermitian.

So, here is an example of pennylane-lightning in action with adjoint differentiation and the object you want calculated!

import pennylane as qml
from pennylane import numpy as np

dev_lightning = qml.device('lightning.qubit', wires=2)
dev = qml.device('default.qubit', wires=2)

@qml.qnode(dev_lightning, diff_method="adjoint")
def circuit_adjoint(a):
    qml.RX(a[0], wires=0)
    qml.CNOT(wires=(0,1))
    qml.RY(a[1], wires=1)
    qml.RZ(a[2], wires=1)
    return qml.expval(qml.PauliX(wires=1))

x = np.array([0.1, 0.2, 0.3], requires_grad=True)
print(qml.jacobian(qml.grad(circuit_adjoint))(x))

I’m not sure if you actually meant to calculate qml.jacobian(qml.grad(circuit_adjoint))(x). In this example, the output seems to be independent of the input. Nevertheless, it runs! If you meant to just calculate, say, the Jacobian, it works as well :).

1 Like

Hi @isaacdevlugt , thanks for your support. I slightly modified your example and shown the output below,

import pennylane as qml
from pennylane import numpy as np
dev_lightning = qml.device('lightning.qubit', wires=2)

@qml.qnode(dev_lightning, diff_method="adjoint", max_diff=2)
def circuit_adjoint(a):
    qml.RX(a[0], wires=0)
    qml.RZ(a[1], wires=1)
    qml.CNOT(wires=(0,1))
    return qml.expval(qml.PauliZ(wires=0))

x = np.array([0.1, 0.2], requires_grad=True)

image

qml.jacobian(qml.grad(circuit_adjoint))(x)
array([[0., 0.],
       [0., 0.]])
qml.gradients.param_shift_hessian(circuit_adjoint)(x)
tensor([[-9.95004165e-01, -8.32667268e-17],
        [-8.32667268e-17,  1.11022302e-16]], requires_grad=True)

Now change the diff_method to ‘best’ and rerun the same lines.

qml.jacobian(qml.grad(circuit_adjoint))(x)
array([[-9.95004165e-01, -8.32667268e-17],
       [-8.32667268e-17, -1.11022302e-16]])

Please compare with above output from parameter shift hessian.

By the way, do you have some idea about dask parallelizing. I have posted something there. It would be great if you can share your views.

Thanks :smile:

Hi @Abhishek
Just jumping in for a comment. The adjoint diff method in lightning.qubit only supports evaluation of Jacobian expectation values. If you would like to evaluate the Hessian, it is best to use the parameter shift method you have already discovered, since param-shift will have all the necessary backend support for the higher-order derivative.

Hope this helps.

Hi @mlxd … Thanks for your support.

Hi, I had a follow up question for when you have multiple inputs with parameter-shift in separate arrays as opposed to 1 array for all inputs:

import pennylane as qml
from pennylane import numpy as np

dev_lightning = qml.device('lightning.qubit', wires=2)
dev = qml.device('default.qubit', wires=2)

@qml.qnode(dev_lightning, "parameter-shift",)
def circuit_adjoint(a,b,c):
    qml.RX(a, wires=0)
    qml.CNOT(wires=(0,1))
    qml.RY(b, wires=1)
    qml.RZ(c, wires=1)
    return qml.expval(qml.PauliX(wires=1))

x_a = np.array([0.1], requires_grad=True)
x_b = np.array([0.2], requires_grad=True)
x_c = np.array([0.3], requires_grad=True)
print(qml.jacobian(qml.grad(circuit_adjoint))(x_a, x_b, x_c))

The above code is not accepted as shown below - how do I deal with multiple inputs?

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
File /quantum_risk_engine/.venv/lib/python3.11/site-packages/numpy/core/fromnumeric.py:57, in _wrapfunc(obj, method, *args, **kwds)
     56 try:
---> 57     return bound(*args, **kwds)
     58 except TypeError:
     59     # A TypeError occurs if the object does have such a method in its
     60     # class, but its signature is not identical to that of NumPy's. This
   (...)
     64     # Call _wrapit from within the except clause to ensure a potential
     65     # exception has a traceback chain.

TypeError: 'ArrayVSpace' object cannot be interpreted as an integer

During handling of the above exception, another exception occurred:

TypeError                                 Traceback (most recent call last)
File /quantum_risk_engine/.venv/lib/python3.11/site-packages/pennylane/_grad.py:330, in jacobian.<locals>._jacobian_function(*args, **kwargs)
    329 try:
--> 330     jac = tuple(_jacobian(func, arg)(*args, **kwargs) for arg in _argnum)
    331 except TypeError as e:

File /quantum_risk_engine/.venv/lib/python3.11/site-packages/pennylane/_grad.py:330, in <genexpr>(.0)
    329 try:
--> 330     jac = tuple(_jacobian(func, arg)(*args, **kwargs) for arg in _argnum)
    331 except TypeError as e:

File /quantum_risk_engine/.venv/lib/python3.11/site-packages/autograd/wrap_util.py:20, in unary_to_nary.<locals>.nary_operator.<locals>.nary_f(*args, **kwargs)
     19     x = tuple(args[i] for i in argnum)
---> 20 return unary_operator(unary_f, x, *nary_op_args, **nary_op_kwargs)

File /quantum_risk_engine/.venv/lib/python3.11/site-packages/autograd/differential_operators.py:64, in jacobian(fun, x)
     63 grads = map(vjp, ans_vspace.standard_basis())
---> 64 return np.reshape(np.stack(grads), jacobian_shape)

File /quantum_risk_engine/.venv/lib/python3.11/site-packages/autograd/tracer.py:48, in primitive.<locals>.f_wrapped(*args, **kwargs)
     47 else:
---> 48     return f_raw(*args, **kwargs)

File <__array_function__ internals>:180, in reshape(*args, **kwargs)

File /quantum_risk_engine/.venv/lib/python3.11/site-packages/numpy/core/fromnumeric.py:298, in reshape(a, newshape, order)
    200 """
    201 Gives a new shape to an array without changing its data.
    202 
   (...)
    296        [5, 6]])
    297 """
--> 298 return _wrapfunc(a, 'reshape', newshape, order=order)

File /quantum_risk_engine/.venv/lib/python3.11/site-packages/numpy/core/fromnumeric.py:66, in _wrapfunc(obj, method, *args, **kwds)
     58 except TypeError:
     59     # A TypeError occurs if the object does have such a method in its
     60     # class, but its signature is not identical to that of NumPy's. This
   (...)
     64     # Call _wrapit from within the except clause to ensure a potential
     65     # exception has a traceback chain.
---> 66     return _wrapit(obj, method, *args, **kwds)

File /quantum_risk_engine/.venv/lib/python3.11/site-packages/numpy/core/fromnumeric.py:43, in _wrapit(obj, method, *args, **kwds)
     42     wrap = None
---> 43 result = getattr(asarray(obj), method)(*args, **kwds)
     44 if wrap:

TypeError: 'ArrayVSpace' object cannot be interpreted as an integer

The above exception was the direct cause of the following exception:

ValueError                                Traceback (most recent call last)
Cell In[12], line 18
     16 x_b = np.array([0.2], requires_grad=True)
     17 x_c = np.array([0.3], requires_grad=True)
---> 18 print(qml.jacobian(qml.grad(circuit_adjoint))(x_a, x_b, x_c))

File /quantum_risk_engine/.venv/lib/python3.11/site-packages/pennylane/_grad.py:333, in jacobian.<locals>._jacobian_function(*args, **kwargs)
    331 except TypeError as e:
    332     if active_return():
--> 333         raise ValueError(
    334             "PennyLane has a new return shape specification that"
    335             " may not work well with autograd and more than one measurement. That may"
    336             " be the source of the error. \n\n"
    337             "See the documentation here for more information:\n"
    338             "https://docs.pennylane.ai/en/stable/introduction/returns.html"
    339         ) from e
    340     raise e
    342 return jac[0] if unpack else jac

ValueError: PennyLane has a new return shape specification that may not work well with autograd and more than one measurement. That may be the source of the error. 

See the documentation here for more information:
https://docs.pennylane.ai/en/stable/introduction/returns.html

Hey @Nikhil_Narayanan!

Your issue is similar to what’s going on here :point_down:

As you can see, there’s a “Wont Fix” label on that :sweat_smile:. However, there are workarounds:

Side note: If doing higher order parameter shift derivatives, then the QNode must specify max_diff=2 .

  1. switch to a single argument:
@qml.qnode(dev_lightning, diff_method="parameter-shift", max_diff=2)
def circuit_adjoint(a):
    qml.RX(a[0], wires=0)
    qml.CNOT(wires=(0, 1))
    qml.RY(a[1], wires=1)
    qml.RZ(a[2], wires=1)
    return qml.expval(qml.PauliX(wires=1))
  1. Use qml.math.stack
dev_lightning = qml.device("lightning.qubit", wires=2)

@qml.qnode(dev_lightning, diff_method="parameter-shift", max_diff=2)
def circuit_adjoint(a, b, c):
    qml.RX(a, wires=0)
    qml.CNOT(wires=(0, 1))
    qml.RY(b, wires=1)
    qml.RZ(c, wires=1)
    return qml.expval(qml.PauliX(wires=1))

x_a = np.array(0.1, requires_grad=True)
x_b = np.array(0.2, requires_grad=True)
x_c = np.array(0.3, requires_grad=True)

def g(*args):
    return qml.math.stack(qml.grad(circuit_adjoint)(*args))

qml.jacobian(g)(x_a, x_b, x_c)

Let me know if that helps!

Hi Isaac,

This helps, but do you know how I could replicate the same functionality with param_shift_hessian? I’m struggling and have tried lots of tricks with stack. I’m asking as I have the same problem as in here as my qml.jacobian returns a 0 gradient but I suspect that the param_shift_hessian would give non 0 results like in Abhishek’s example?

qml.gradients.param_shift_hessian(circuit_adjoint)(x) tensor([[-9.95004165e-01, -8.32667268e-17], [-8.32667268e-17, 1.11022302e-16]], requires_grad=True)

In other words, how do I make something like the below work?


dev_lightning = qml.device("lightning.qubit", wires=2)

@qml.qnode(dev_lightning, diff_method="parameter-shift", max_diff=2)
def circuit_adjoint(a, b, c):
    qml.RX(a, wires=0)
    qml.CNOT(wires=(0, 1))
    qml.RY(b, wires=1)
    qml.RZ(c, wires=1)
    return qml.expval(qml.PauliX(wires=1))

x_a = np.array(0.1, requires_grad=True)
x_b = np.array(0.2, requires_grad=True)
x_c = np.array(0.3, requires_grad=True)

def g(*args):
    return qml.gradients.param_shift_hessian((circuit_adjoint)(*args))

g(x_a, x_b, x_c)

In other words, how do I make something like the below work?

I think you have (*args) in the wrong spot. This produces a result that’s not a function:

dev_lightning = qml.device("lightning.qubit", wires=2)

@qml.qnode(dev_lightning, diff_method="parameter-shift", max_diff=2)
def circuit_adjoint(a, b, c):
    qml.RX(a, wires=0)
    qml.CNOT(wires=(0, 1))
    qml.RY(b, wires=1)
    qml.RZ(c, wires=1)
    return qml.expval(qml.PauliX(wires=1))

x_a = np.array(0.1, requires_grad=True)
x_b = np.array(0.2, requires_grad=True)
x_c = np.array(0.3, requires_grad=True)

def g(*args):
    return qml.gradients.param_shift_hessian((circuit_adjoint))(*args)

g(x_a, x_b, x_c)
(-0.18884787122725294, -0.18884787122240193, -0.1888478712274544)

qml.jacobian returns a 0 gradient

Hmm, I can’t replicate that :thinking:. This is what I get:

dev_lightning = qml.device("lightning.qubit", wires=2)


@qml.qnode(dev_lightning, diff_method="parameter-shift", max_diff=2)
def circuit_adjoint(a, b, c):
    qml.RX(a, wires=0)
    qml.CNOT(wires=(0, 1))
    qml.RY(b, wires=1)
    qml.RZ(c, wires=1)
    return qml.expval(qml.PauliX(wires=1))


x_a = np.array(0.1, requires_grad=True)
x_b = np.array(0.2, requires_grad=True)
x_c = np.array(0.3, requires_grad=True)


def g(*args):
    return qml.math.stack(qml.grad(circuit_adjoint)(*args))


qml.jacobian(g)(x_a, x_b, x_c)
(array([-0.18884787, -0.09347337,  0.0058613 ]),
 array([-0.09347337, -0.18884787, -0.28818254]),
 array([ 0.0058613 , -0.28818254, -0.18884787]))

Let me know if this helps! Here’s my qml.about() in case there’s an inconsistency there:

Name: PennyLane
Version: 0.31.0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/PennyLaneAI/pennylane

Platform info:           macOS-13.4.1-x86_64-i386-64bit
Python version:          3.9.14
Numpy version:           1.23.5
Scipy version:           1.10.0
Installed devices:
- default.gaussian (PennyLane-0.31.0)
- default.mixed (PennyLane-0.31.0)
- default.qubit (PennyLane-0.31.0)
- default.qubit.autograd (PennyLane-0.31.0)
- default.qubit.jax (PennyLane-0.31.0)
- default.qubit.tf (PennyLane-0.31.0)
- default.qubit.torch (PennyLane-0.31.0)
- default.qutrit (PennyLane-0.31.0)
- null.qubit (PennyLane-0.31.0)
- lightning.qubit (PennyLane-Lightning-0.31.0)