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:
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1436?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1436](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1436?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (c6dcd35) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/614db57b02536bb09fbfbae853ceff8878b928af?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (614db57) will **not change** coverage.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1436/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1436?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@           Coverage Diff           @@
##           master    #1436   +/-   ##
=======================================
  Coverage   98.23%   98.23%           
=======================================
  Files         160      160           
  Lines       11966    11966           
=======================================
  Hits        11755    11755           
  Misses        211      211           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1436?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/interfaces/tf.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1436/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvdGYucHk=) | `97.43% <ø> (ø)` | |
| [pennylane/interfaces/torch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1436/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvdG9yY2gucHk=) | `100.00% <ø> (ø)` | |
| [pennylane/interfaces/autograd.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1436/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYXV0b2dyYWQucHk=) | `100.00% <100.00%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1436?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1436?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [614db57...c6dcd35](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1436?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).


Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index 211983c0a..ec66a218a 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -43,7 +43,12 @@
 
 * Quantum function transforms now preserve the format of the measurement
   results, so that a single measurement returns a single value rather than
-  an array with a single element. [(#1434)](https://github.com/PennyLaneAI/pennylane/pull/1434/files)
+  an array with a single element. [(#1434)](https://github.com/PennyLaneAI/pennylane/pull/1434)
+
+* Fixed a bug in the parameter-shift Hessian implementation, which resulted
+  in the incorrect Hessian being returned for a cost function
+  that performed post-processing on a vector-valued QNode.
+  [(#)](https://github.com/PennyLaneAI/pennylane/pull/)
 
 <h3>Documentation</h3>
 
@@ -51,7 +56,7 @@
 
 This release contains contributions from (in alphabetical order):
 
-Olivia Di Matteo, Josh Izaac, Ashish Panigrahi
+Olivia Di Matteo, Josh Izaac, Ashish Panigrahi.
 
 
 # Release 0.16.0 (current release)
diff --git a/pennylane/interfaces/autograd.py b/pennylane/interfaces/autograd.py
index 80211acc5..5edd5bf15 100644
--- a/pennylane/interfaces/autograd.py
+++ b/pennylane/interfaces/autograd.py
@@ -259,7 +259,7 @@ class AutogradInterface(AnnotatedQueue):
                     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
index 0688dbda6..9c494c362 100644
--- a/pennylane/interfaces/tf.py
+++ b/pennylane/interfaces/tf.py
@@ -212,7 +212,11 @@ class TFInterface(AnnotatedQueue):
 
                     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
index b8a1dfd59..2bfdfcbcb 100644
--- 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 @@ class _TorchInterface(torch.autograd.Function):
 
                 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
index 02d5de08f..df8f94e51 100644
--- a/tests/interfaces/test_qnode_autograd.py
+++ b/tests/interfaces/test_qnode_autograd.py
@@ -917,6 +917,60 @@ class TestQNode:
         ]
         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
index 114d1bd1f..c674c3b5c 100644
--- a/tests/interfaces/test_qnode_tf.py
+++ b/tests/interfaces/test_qnode_tf.py
@@ -722,6 +722,58 @@ class TestQNode:
 
         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
index fc111a9c9..e227f9e41 100644
--- a/tests/interfaces/test_qnode_torch.py
+++ b/tests/interfaces/test_qnode_torch.py
@@ -735,6 +735,60 @@ class TestQNode:
         ]
         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):