Repository: PennyLaneAI/pennylane

Issue:
Parameter-shift Hessian of multi-observable circuits incorrect
#### Issue description

Consider a circuit with an iterable of `qml.expval`s returned and a cost function that sums this iterable (using the PennyLane version of numpy):

```python
dev = qml.device('default.qubit', wires=1)

@qml.qnode(dev, diff_method='parameter-shift')
def _circ(param):
    qml.RX(param[0], wires=0)
    return [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(0))]

circ = lambda p: np.sum(_circ(p))
```
Now evaluate the second order derivative of this circuit, once via manual finite difference, once via parameter shifts:
```python
param = np.random.random(1)
delta = 1e-5
print(((circ(param+delta)-circ(param))/delta-(circ(param)-circ(param-delta))/delta)/delta)
>>> -1.957101147809226
print(qml.jacobian(qml.jacobian(circ))(param))
>>> -3.91459319
```
We consistently find the latter output to be twice the correct value. When switching to `diff_method='backprop'`, the finite difference and the `qml.jacobian(qml.grad(circ))` results coincide.

For larger parameter spaces, the full Hessian is off by the same factor `2`, i.e. both off-diagonal and diagonal entries are affected likewise.

* *Expected behavior:* (What you expect to happen)
Hessian computed via parameter shifts should match the one computed via backprop/the manually computed value.

* *Actual behavior:* (What actually happens)
The parameter shift Hessian is twice as large as the exact value.

* *Reproduces how often:* (What percentage of the time does it reproduce?)
Always

* *System information:* (post the output of `import pennylane as qml; qml.about()`)

```python
Name: PennyLane
Version: 0.15.1
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/XanaduAI/pennylane
Author: None
Author-email: None
License: Apache License 2.0
Location: /home/david/venvs/xanadu/lib/python3.8/site-packages
Requires: autograd, numpy, scipy, networkx, appdirs, semantic-version, toml
Required-by: PennyLane-SF, pennylane-qulacs, PennyLane-qiskit, PennyLane-Qchem, PennyLane-Forest, PennyLane-Cirq
Platform info:           Linux-5.4.0-74-generic-x86_64-with-glibc2.29
Python version:          3.8.5
Numpy version:           1.18.5
Scipy version:           1.6.3
Installed devices:
- default.gaussian (PennyLane-0.15.1)
- default.mixed (PennyLane-0.15.1)
- default.qubit (PennyLane-0.15.1)
- default.qubit.autograd (PennyLane-0.15.1)
- default.qubit.jax (PennyLane-0.15.1)
- default.qubit.tf (PennyLane-0.15.1)
- default.tensor (PennyLane-0.15.1)
- default.tensor.tf (PennyLane-0.15.1)
- strawberryfields.fock (PennyLane-SF-0.15.0)
- strawberryfields.gaussian (PennyLane-SF-0.15.0)
- strawberryfields.gbs (PennyLane-SF-0.15.0)
- strawberryfields.remote (PennyLane-SF-0.15.0)
- strawberryfields.tf (PennyLane-SF-0.15.0)
- qulacs.simulator (pennylane-qulacs-0.15.0)
- qiskit.aer (PennyLane-qiskit-0.15.0)
- qiskit.basicaer (PennyLane-qiskit-0.15.0)
- qiskit.ibmq (PennyLane-qiskit-0.15.0)
- forest.numpy_wavefunction (PennyLane-Forest-0.15.0)
- forest.qvm (PennyLane-Forest-0.15.0)
- forest.wavefunction (PennyLane-Forest-0.15.0)
- cirq.mixedsimulator (PennyLane-Cirq-0.15.0)
- cirq.pasqal (PennyLane-Cirq-0.15.0)
- cirq.qsim (PennyLane-Cirq-0.15.0)
- cirq.qsimh (PennyLane-Cirq-0.15.0)
- cirq.simulator (PennyLane-Cirq-0.15.0)
```



Discussion:
Taking a look.
Pretty sure the problem is on [interfaces/autograd.py Line 262](https://github.com/PennyLaneAI/pennylane/blob/25284ed58880fa6658676d2184eff5dd4503d699/pennylane/interfaces/autograd.py#L262).

But I never quite understood that line in the first place, so I'll try and work out what that line is supposed to be.
Ah, I can provide some insight, but not all unfortunately.

* In the case that the QNode simply returns a scalar, then `dy == [1]`, and the vector-Hessian product is simply `ddy @ hessian`.

* If the QNode instead returns a _vector_, then this section of the code will be called twice. The first time with `dy = [1, 0]` (I believe), and the second time with `dy = [0, 1]`. Hence why we cache the Hessian matrix and contract with `dy`: `vhp = dy @ ddy @ hessian @ dy.T`.

I wonder if this line is missing a normalization factor? We have tests where this line passes, so maybe the normalization is `dy`/`ddy`/`ans` dependent?

Resolving patch:
diff --git a/pennylane/interfaces/autograd.py b/pennylane/interfaces/autograd.py
--- a/pennylane/interfaces/autograd.py
+++ b/pennylane/interfaces/autograd.py
@@ -259,7 +259,7 @@ def hessian_product(ddy):
                     hessian = _evaluate_grad_matrix(p, "hessian")
 
                     if dy.size > 1:
-                        vhp = dy @ ddy @ hessian @ dy.T
+                        vhp = dy @ ddy @ hessian @ dy.T / np.linalg.norm(dy) ** 2
                     else:
                         vhp = ddy @ hessian
                         vhp = vhp.flatten()
diff --git a/pennylane/interfaces/tf.py b/pennylane/interfaces/tf.py
--- a/pennylane/interfaces/tf.py
+++ b/pennylane/interfaces/tf.py
@@ -212,7 +212,11 @@ def hessian_product(ddy, **tfkwargs):
 
                     vhp = tf.cond(
                         tf.rank(hessian) > 2,
-                        lambda: dy_row @ ddy @ hessian @ tf.transpose(dy_row),
+                        lambda: dy_row
+                        @ ddy
+                        @ hessian
+                        @ tf.transpose(dy_row)
+                        / tf.linalg.norm(dy_row) ** 2,
                         lambda: ddy @ hessian,
                     )
 
diff --git a/pennylane/interfaces/torch.py b/pennylane/interfaces/torch.py
--- a/pennylane/interfaces/torch.py
+++ b/pennylane/interfaces/torch.py
@@ -15,7 +15,7 @@
 This module contains the mixin interface class for creating differentiable quantum tapes with
 PyTorch.
 """
-# pylint: disable=protected-access, attribute-defined-outside-init, arguments-differ, no-member, import-self
+# pylint: disable=protected-access, attribute-defined-outside-init, arguments-differ, no-member, import-self, too-many-statements
 import numpy as np
 import semantic_version
 import torch
@@ -144,6 +144,7 @@ def backward(ctx_, ddy):  # pragma: no cover
 
                 if torch.squeeze(ddy).ndim > 1:
                     vhp = ctx_.dy.view(1, -1) @ ddy @ hessian @ ctx_.dy.view(-1, 1)
+                    vhp = vhp / torch.linalg.norm(ctx_.dy) ** 2
                 else:
                     vhp = ddy @ hessian
 


Test patch:
diff --git a/tests/interfaces/test_qnode_autograd.py b/tests/interfaces/test_qnode_autograd.py
--- a/tests/interfaces/test_qnode_autograd.py
+++ b/tests/interfaces/test_qnode_autograd.py
@@ -917,6 +917,60 @@ def circuit(x):
         ]
         assert np.allclose(hess, expected_hess, atol=tol, rtol=0)
 
+    def test_hessian_vector_valued_postprocessing(self, dev_name, diff_method, mocker, tol):
+        """Test hessian calculation of a vector valued QNode with post-processing"""
+        if diff_method not in {"parameter-shift", "backprop"}:
+            pytest.skip("Test only supports parameter-shift or backprop")
+
+        dev = qml.device(dev_name, wires=1)
+
+        @qnode(dev, diff_method=diff_method, interface="autograd")
+        def circuit(x):
+            qml.RX(x[0], wires=0)
+            qml.RY(x[1], wires=0)
+            return [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(0))]
+
+        def cost_fn(x):
+            return x @ circuit(x)
+
+        x = np.array([0.76, -0.87], requires_grad=True)
+        res = cost_fn(x)
+
+        a, b = x
+
+        expected_res = x @ [np.cos(a) * np.cos(b), np.cos(a) * np.cos(b)]
+        assert np.allclose(res, expected_res, atol=tol, rtol=0)
+
+        grad_fn = qml.grad(cost_fn)
+        g = grad_fn(x)
+
+        expected_g = [
+            np.cos(b) * (np.cos(a) - (a + b) * np.sin(a)),
+            np.cos(a) * (np.cos(b) - (a + b) * np.sin(b)),
+        ]
+        assert np.allclose(g, expected_g, atol=tol, rtol=0)
+
+        spy = mocker.spy(JacobianTape, "hessian")
+        hess = qml.jacobian(grad_fn)(x)
+
+        if diff_method == "backprop":
+            spy.assert_not_called()
+        elif diff_method == "parameter-shift":
+            spy.assert_called_once()
+
+        expected_hess = [
+            [
+                -(np.cos(b) * ((a + b) * np.cos(a) + 2 * np.sin(a))),
+                -(np.cos(b) * np.sin(a)) + (-np.cos(a) + (a + b) * np.sin(a)) * np.sin(b),
+            ],
+            [
+                -(np.cos(b) * np.sin(a)) + (-np.cos(a) + (a + b) * np.sin(a)) * np.sin(b),
+                -(np.cos(a) * ((a + b) * np.cos(b) + 2 * np.sin(b))),
+            ],
+        ]
+
+        assert np.allclose(hess, expected_hess, atol=tol, rtol=0)
+
     def test_hessian_vector_valued_separate_args(self, dev_name, diff_method, mocker, tol):
         """Test hessian calculation of a vector valued QNode that has separate input arguments"""
         if diff_method not in {"parameter-shift", "backprop"}:
diff --git a/tests/interfaces/test_qnode_tf.py b/tests/interfaces/test_qnode_tf.py
--- a/tests/interfaces/test_qnode_tf.py
+++ b/tests/interfaces/test_qnode_tf.py
@@ -722,6 +722,58 @@ def circuit(x):
 
         np.testing.assert_allclose(hess, expected_hess, atol=tol, rtol=0, verbose=True)
 
+    def test_hessian_vector_valued_postprocessing(self, dev_name, diff_method, mocker, tol):
+        """Test hessian calculation of a vector valued QNode with post-processing"""
+        if diff_method not in {"parameter-shift", "backprop"}:
+            pytest.skip("Test only supports parameter-shift or backprop")
+
+        dev = qml.device(dev_name, wires=1)
+
+        @qnode(dev, diff_method=diff_method, interface="tf")
+        def circuit(x):
+            qml.RX(x[0], wires=0)
+            qml.RY(x[1], wires=0)
+            return [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(0))]
+
+        x = tf.Variable([0.76, -0.87], dtype=tf.float64)
+
+        with tf.GradientTape(persistent=True) as tape1:
+            with tf.GradientTape(persistent=True) as tape2:
+                res = tf.tensordot(x, circuit(x), axes=[0, 0])
+
+            spy = mocker.spy(JacobianTape, "hessian")
+            g = tape2.jacobian(res, x, experimental_use_pfor=False)
+
+        hess = tape1.jacobian(g, x, experimental_use_pfor=False)
+
+        if diff_method == "parameter-shift":
+            spy.assert_called_once()
+        elif diff_method == "backprop":
+            spy.assert_not_called()
+
+        a, b = x * 1.0
+
+        expected_res = a * tf.cos(a) * tf.cos(b) + b * tf.cos(a) * tf.cos(b)
+        assert np.allclose(res, expected_res, atol=tol, rtol=0)
+
+        expected_g = [
+            tf.cos(b) * (tf.cos(a) - (a + b) * tf.sin(a)),
+            tf.cos(a) * (tf.cos(b) - (a + b) * tf.sin(b)),
+        ]
+        assert np.allclose(g, expected_g, atol=tol, rtol=0)
+
+        expected_hess = [
+            [
+                -(tf.cos(b) * ((a + b) * tf.cos(a) + 2 * tf.sin(a))),
+                -(tf.cos(b) * tf.sin(a)) + (-tf.cos(a) + (a + b) * tf.sin(a)) * tf.sin(b),
+            ],
+            [
+                -(tf.cos(b) * tf.sin(a)) + (-tf.cos(a) + (a + b) * tf.sin(a)) * tf.sin(b),
+                -(tf.cos(a) * ((a + b) * tf.cos(b) + 2 * tf.sin(b))),
+            ],
+        ]
+        assert np.allclose(hess, expected_hess, atol=tol, rtol=0)
+
     def test_hessian_ragged(self, dev_name, diff_method, mocker, tol):
         """Test hessian calculation of a ragged QNode"""
         if diff_method not in {"parameter-shift", "backprop"}:
diff --git a/tests/interfaces/test_qnode_torch.py b/tests/interfaces/test_qnode_torch.py
--- a/tests/interfaces/test_qnode_torch.py
+++ b/tests/interfaces/test_qnode_torch.py
@@ -735,6 +735,60 @@ def circuit(x):
         ]
         assert np.allclose(hess.detach(), expected_hess, atol=tol, rtol=0)
 
+    def test_hessian_vector_valued_postprocessing(self, dev_name, diff_method, mocker, tol):
+        """Test hessian calculation of a vector valued QNode with post-processing"""
+        if diff_method not in {"parameter-shift", "backprop"}:
+            pytest.skip("Test only supports parameter-shift or backprop")
+
+        dev = qml.device(dev_name, wires=1)
+
+        @qnode(dev, diff_method=diff_method, interface="torch")
+        def circuit(x):
+            qml.RX(x[0], wires=0)
+            qml.RY(x[1], wires=0)
+            return [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(0))]
+
+        x = torch.tensor([0.76, -0.87], requires_grad=True, dtype=torch.float64)
+
+        def cost_fn(x):
+            return x @ circuit(x)
+
+        a, b = x.detach().numpy()
+
+        res = cost_fn(x)
+        expected_res = np.array([a, b]) @ [np.cos(a) * np.cos(b), np.cos(a) * np.cos(b)]
+        assert np.allclose(res.detach(), expected_res, atol=tol, rtol=0)
+
+        res.backward()
+
+        g = x.grad
+        expected_g = [
+            np.cos(b) * (np.cos(a) - (a + b) * np.sin(a)),
+            np.cos(a) * (np.cos(b) - (a + b) * np.sin(b)),
+        ]
+        assert np.allclose(g.detach(), expected_g, atol=tol, rtol=0)
+
+        spy = mocker.spy(JacobianTape, "hessian")
+        hess = hessian(cost_fn, x)
+
+        if diff_method == "backprop":
+            spy.assert_not_called()
+        elif diff_method == "parameter-shift":
+            spy.assert_called_once()
+
+        expected_hess = [
+            [
+                -(np.cos(b) * ((a + b) * np.cos(a) + 2 * np.sin(a))),
+                -(np.cos(b) * np.sin(a)) + (-np.cos(a) + (a + b) * np.sin(a)) * np.sin(b),
+            ],
+            [
+                -(np.cos(b) * np.sin(a)) + (-np.cos(a) + (a + b) * np.sin(a)) * np.sin(b),
+                -(np.cos(a) * ((a + b) * np.cos(b) + 2 * np.sin(b))),
+            ],
+        ]
+
+        assert np.allclose(hess.detach(), expected_hess, atol=tol, rtol=0)
+
 
 class Test_adjoint:
     def test_adjoint_default_save_state(self, mocker):