Support classical fisher gradients when using Autograd

[josh146]

Context: The qml.math.divide() function was producing a list of a list of ArrayBox objects when using autograd, which was causing differentiation to fail. It seems that either swapping back to standard python division, or stacking the output of qml.math.divide(), fixes the problem.

Description of the Change:

• Swapped back to standard Python division, since this is likely faster than adding additional stacks.

Benefits: differentiating the fisher information now works:

dev = qml.device("default.qubit", wires=3)

@qml.qnode(dev)
def circuit(weights):
    qml.RX(weights[0], wires=0)
    qml.RX(weights[1], wires=1)
    qml.RX(weights[2], wires=2)
    qml.CNOT(wires=[0, 1])
    qml.RY(weights[4], wires=0)
    qml.RY(weights[5], wires=1)
    qml.RY(weights[6], wires=2)
    qml.CNOT(wires=[1, 2])
    qml.CRX(weights[7], wires=[0, 1])
    qml.CRY(weights[8], wires=[1, 2])
    qml.CRZ(weights[9], wires=[2, 0])
    return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))

weights = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0], requires_grad=True)
print(qml.jacobian(qml.qinfo.classical_fisher(circuit), argnum=0)(weights))

Possible Drawbacks: I haven't verified if the gradient is correct or not 🙂

Related GitHub Issues: n/a

[codecov]

Codecov Report

❗ No coverage uploaded for pull request base (v0.24.0-rc0@76e3fa3). Click here to learn what that means.
The diff coverage is n/a.

@@              Coverage Diff               @@
##             v0.24.0-rc0    #2688   +/-   ##
==============================================
  Coverage               ?   99.61%           
==============================================
  Files                  ?      251           
  Lines                  ?    20670           
  Branches               ?        0           
==============================================
  Hits                   ?    20590           
  Misses                 ?       80           
  Partials               ?        0           

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 76e3fa3...084d4f1. Read the comment docs.

[Qottmann]

This PR is now based on master, but I guess we would want this to be in the new release candidate v0.24.0-rc0? Can I just switch the base or will I wreak havoc by doing this?

I.e. I would like to add a differentiability test for autograd before merging.

[Qottmann]

Jax and autograd now produce the correct result for the simple analytic example. Tensorflow and torch still yield the wrong result. Here are the two examples that are not working. Would be worth adding this to the doc string, warning the user that differentiating with torch and tf does not work atm.

import torch

@qml.qnode(dev, interface="torch")
def circ(params):
    qml.RY(params, wires=0)
    return qml.probs(wires=range(1))

params = np.pi / 4 * torch.tensor([1.0], requires_grad=True) # note that this is arbitrary, result is constant for any value

print("Check analytical result: ", np.allclose(qml.qinfo.classical_fisher(circ)(params).detach().numpy(),1))

result = np.zeros((1, 1, 1), dtype="float64")

result_calc = torch.autograd.functional.jacobian(qml.qinfo.classical_fisher(circ), params)

print("Check analytic gradient: ", np.allclose(result, result_calc, atol=1e-6), "value = ", result_calc)

Check analytical result:  True
Check analytic gradient:  False value =  tensor([[[-2.0000]]])

import tensorflow as tf

@qml.qnode(dev, interface="tf")
def circ(params):
    qml.RY(params, wires=0)
    return qml.probs(wires=range(1))

params = tf.Variable([np.pi/4])

print("Check analytical result: ", np.allclose(qml.qinfo.classical_fisher(circ)(params),1))

result = np.zeros((1, 1, 1), dtype="float64")

with tf.GradientTape() as tape:
    loss = qml.qinfo.classical_fisher(circ)(params)
result_calc = tape.jacobian(loss, params)

print("Check analytic gradient: ", np.allclose(result, result_calc, atol=1e-6), "value = ", result_calc)

Check analytical result:  True
Check analytic gradient:  False value =  tf.Tensor([[[-1.9999996]]], shape=(1, 1, 1), dtype=float32)

[Qottmann]

Some further testing, this time no analytic results but just for consistency. It seems that (torch and tf) and (autograd and jax) are producing consistent results.

dev = qml.device("default.qubit", wires=3)

def qfunc(weights):
    qml.RX(weights[0], wires=0)
    qml.RX(weights[1], wires=1)
    qml.RX(weights[2], wires=2)
    qml.CNOT(wires=[0, 1])
    qml.RY(weights[4], wires=0)
    qml.RY(weights[5], wires=1)
    qml.RY(weights[6], wires=2)
    qml.CNOT(wires=[1, 2])
    qml.CRX(weights[7], wires=[0, 1])
    qml.CRY(weights[8], wires=[1, 2])
    qml.CRZ(weights[9], wires=[2, 0])
    return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))

# Compute gradients of CFIM for different interfaces
circuit = qml.QNode(qfunc, dev, interface="torch")
weights = torch.tensor([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0], requires_grad=True)
grad_torch = torch.autograd.functional.jacobian(qml.qinfo.classical_fisher(circuit), weights)

circuit = qml.QNode(qfunc, dev, interface="autograd")
weights = pnp.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0], requires_grad=True)
grad_autograd = qml.jacobian(qml.qinfo.classical_fisher(circuit))(weights)

circuit = qml.QNode(qfunc, dev, interface="jax")
weights = jnp.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0])
grad_jax = jax.jacobian(qml.qinfo.classical_fisher(circuit))(weights)

circuit = qml.QNode(qfunc, dev, interface="tf")
weights = tf.Variable([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0])
with tf.GradientTape() as tape:
    loss = qml.qinfo.classical_fisher(circuit)(weights)
grad_tf = tape.jacobian(loss, weights)

# Evaluate and compare
grads = [grad_autograd, grad_tf, grad_torch, grad_jax]

is_same = np.zeros((4,4))

for i, g1 in enumerate(grads):
    for j, g2 in enumerate(grads):
        is_same[i, j] = np.allclose(g1, g2, atol=1e-6)

print(is_same)

[[1. 0. 0. 1.]
 [0. 1. 1. 0.]
 [0. 1. 1. 0.]
 [1. 0. 0. 1.]]

[josh146]

Thanks for checking this @Qottmann!

Regarding the TF and Torch examples; it might be really good to add Torch and TF tests to the test suite for the gradient (alongside JAX and Autograd gradient tests), but simply mark the Torch and TF tests as @pytest.mark.xfail. This way, pytest expects the test to fail.