Repository: PennyLaneAI/pennylane (python)

Issue:
Problems with the parameter shift hessian
The parameter shift Hessian does not seem to work reliably in all cases. It is difficult to come up with a very minimal example demonstrating the problem as I do not yet fully understand what what is going wrong. But I managed to condense it down to the following "minimal" example that demonstrates the main bug and shows a few minor problems. Hopefully, e.g.,  @josh146  can take it from there and get to the bottom of this.

The main problem is that the Autograd Hessian of a function involving a parameter-shift QNode does not agree with a finite difference approximation of the Hessian or the "best" method Hessian (the later two do agree to the expected accuracy).

There are a few additional minor issues that can be triggered by modifying the code according to the comments:
The first being an IndexError when a gate is commented from the ansatz, making the result independent of one of the paramters.
The second being a error that occurs when trying to draw a circuit during Hessian computation.    

Here is the example:
```
import pennylane as qml
import numpy as np
import copy
from pennylane.utils import _flatten, unflatten

def finite_diff_hessian(energy, params, shift=1e-4, **kwargs):
    """Computes a finite difference approximation to the Hessian"""
    flat_params = copy.deepcopy(np.array(list(_flatten(params))))
    hessian = np.zeros((len(flat_params), len(flat_params)))
    for idx1 in range(len(flat_params)):
        params_shift1 = np.array([shift if idx==idx1 else 0. for idx in range(len(flat_params))])
        for idx2 in range(len(flat_params)):
            params_shift2 = np.array([shift if idx==idx2 else 0. for idx in range(len(flat_params))])

            energy11 = energy(unflatten(flat_params+params_shift1+params_shift2, params))
            energy01 = energy(unflatten(flat_params+params_shift1, params))
            energy10 = energy(unflatten(flat_params+params_shift2, params))
            energy00 = energy(params)
            hessian[idx1, idx2] = (energy11 - energy01 - energy10 + energy00)/(shift**2)
    #hessian = 0.5*(hessian + hessian.T) # symmetrization is not necssary
    return hessian

nwires = 8
wires = range(nwires)
dev = qml.device('default.qubit', wires=nwires)

def make_cost(diff_method='parameter-shift'):
    @qml.template
    def ansatz(params, wires=wires):
        qml.RY(1.5, wires=[6])
        qml.CNOT(wires=[6, 4])
        qml.CNOT(wires=[6, 7])
        qml.CNOT(wires=[4, 5])
        qml.SWAP(wires=[2, 4])
        qml.SWAP(wires=[3, 5])
        qml.CRY(params[0], wires=[2, 4])
        qml.SWAP(wires=[0, 2])
        qml.SWAP(wires=[1, 3])
        qml.CRY(params[1], wires=[0, 2]) # commenting this produces: "IndexError: too many indices for array: array is 0-dimensional, but 1 were indexed"
        

    def run_circuit(params, ob, wires=wires):
        @qml.qnode(dev, interface="autograd", diff_method=diff_method) #without diff_method='parameter-shift' this works!
        def circ(params, wires=wires):
            ansatz(params, wires=wires)
            return qml.expval(ob)
        ret = circ(params, wires=wires)
        #print(circ.draw()) # produces "TypeError: unsupported format string passed to ArrayBox.__format__" Maybe another bug?
        return ret

    def circuit(params, wires=wires):
        return run_circuit(params, qml.PauliX(0) @ qml.PauliX(1) @ qml.PauliZ(2) \
                           @ qml.PauliZ(3) @ qml.PauliZ(4) @ qml.PauliZ(5) @ qml.PauliX(6) @ qml.PauliX(7), wires=wires)
        

    return circuit

def energy(params):
    return make_cost(diff_method='parameter-shift')(params)

def energy2(params):
    return make_cost(diff_method='best')(params)

params = np.array([-0.02, 0.4])

fd_hessian = finite_diff_hessian(energy, params)
pl_hessian = qml.jacobian(qml.grad(energy))(params)
pl_hessian2 = qml.jacobian(qml.grad(energy2))(params)

assert np.allclose(fd_hessian.ravel(), pl_hessian2.ravel(), rtol=0, atol=1e-5), "best method Hessian did not match.\nfd_hessian:\n{}\npl_hessian:\n{}\ndifference:\n{}".format(fd_hessian, pl_hessian, fd_hessian.ravel() - pl_hessian.ravel())
print("best method Hessians did match")
assert np.allclose(fd_hessian.ravel(), pl_hessian.ravel(), rtol=0, atol=1e-5), "parameter-shift Hessian did not match.\nfd_hessian:\n{}\npl_hessian:\n{}\ndifference:\n{}".format(fd_hessian, pl_hessian, fd_hessian.ravel() - pl_hessian.ravel())
print("parameter-shift Hessians did match")
```
The expected output is 
```
best method Hessians did match
parameter-shift Hessians did match
```
instead I am getting
```
best method Hessians did match
AssertionError                            Traceback (most recent call last)
<ipython-input-1-d59f46ef7659> in <module>
     70 assert np.allclose(fd_hessian.ravel(), pl_hessian2.ravel(), rtol=0, atol=1e-5), "best method Hessian did not match.\nfd_hessian:\n{}\npl_hessian:\n{}\ndifference:\n{}".format(fd_hessian, pl_hessian, fd_hessian.ravel() - pl_hessian.ravel())
     71 print("best method Hessians did match")
---> 72 assert np.allclose(fd_hessian.ravel(), pl_hessian.ravel(), rtol=0, atol=1e-5), "parameter-shift Hessian did not match.\nfd_hessian:\n{}\npl_hessian:\n{}\ndifference:\n{}".format(fd_hessian, pl_hessian, fd_hessian.ravel() - pl_hessian.ravel())
     73 print("parameter-shift Hessians did match")

AssertionError: parameter-shift Hessian did not match.
fd_hessian:
[[-0.24439083 -0.00049429]
 [-0.00049429 -0.24438817]]
pl_hessian:
[[-4.83893333e-01 -4.95420897e-04]
 [-4.95420897e-04 -5.87862185e-01]]
difference:
[2.39502503e-01 1.12740152e-06 1.12740152e-06 3.43474020e-01]
```
Maybe this has to do with the non standard grad recipe of CRY? 

Pull request:
Fixes two bugs in the parameter-shift Hessian
**Context:**

As noted in #1258,

* Having an unused parameter in the Autograd interface results in an indexing error during backpropagation when computing the parameter-shift Hessian.

* The parameter-shift Hessian only supports the two-term parameter-shift rule currently, and will return incorrect results for operations not supporting the two-term rule.

In the latter case, the implementation of the parameter-shift Hessian is hardcoded to follow the results from https://arxiv.org/abs/2008.06517. One short-term solution could be to expand the tape around the controlled rotation gates inside the parameter-shift Hessian logic. Unfortunately, the controlled rotation decomposition are not 1-1 and involve classical processing. I attempted to implement this, but the complexity spiraled significantly. Likely the best solution is to:

* For now, raise an error if controlled rotation gates occur
* Suggest to the user they decompose these gates
* Long term, re-write the Hessian logic to be more general.

**Description of the Change:**

- Fixes the first issue by flattening the vector-Hessian product before returning it to Autograd, and adds a test.

- The parameter-shift Hessian now raises an error if an unsupported gate is encountered, and provides a useful error message:

  ```python
  @qml.qnode(dev, diff_method="parameter-shift")
  def circuit(b):
      qml.QubitStateVector(init_state, wires=0)
      qml.CRY(b, wires=[0, 1])
      return qml.expval(qml.PauliX(0))
  ```

  ```pycon
  >>> grad = qml.grad(qml.grad(circuit))(0.123)
  ValueError: The operation CRX is currently not supported for the parameter-shift Hessian.
  Please decompose the operation in your QNode by replacing it with 'CRX.decomposition(0.123, wires=[0, 1])'
  ```

**Benefits:** Fixes the first bug, while raising a more useful error message for the second bug as we work to support that use-case.

**Possible Drawbacks:** n/a

**Related GitHub Issues:** Fixes #1258 


Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index 107b90f7d..9b72ef04a 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -48,6 +48,16 @@
 
 <h3>Bug fixes</h3>
 
+* Fixes two bugs in the parameter-shift Hessian.
+  [(#1260)](https://github.com/PennyLaneAI/pennylane/pull/1260)
+
+  - Fixes a bug where having an unused parameter in the Autograd interface
+    would result in an indexing error during backpropagation.
+
+  - The parameter-shift Hessian only supports the two-term parameter-shift
+    rule currently, so raises an error if asked to differentiate
+    any unsupported gates (such as the controlled rotation gates).
+
 * A bug which resulted in `qml.adjoint()` and `qml.inv()` failing to work with
   templates has been fixed.
   [(#1243)](https://github.com/PennyLaneAI/pennylane/pull/1243)
diff --git a/pennylane/interfaces/autograd.py b/pennylane/interfaces/autograd.py
index 409c29a5b..ec138282f 100644
--- a/pennylane/interfaces/autograd.py
+++ b/pennylane/interfaces/autograd.py
@@ -261,7 +261,8 @@ class AutogradInterface(AnnotatedQueue):
                     if dy.size > 1:
                         vhp = dy @ ddy @ hessian @ dy.T
                     else:
-                        vhp = np.squeeze(ddy @ hessian)
+                        vhp = ddy @ hessian
+                        vhp = vhp.flatten()
 
                     return vhp
 
diff --git a/pennylane/tape/jacobian_tape.py b/pennylane/tape/jacobian_tape.py
index 43cacf874..5cddd1205 100644
--- a/pennylane/tape/jacobian_tape.py
+++ b/pennylane/tape/jacobian_tape.py
@@ -685,6 +685,36 @@ class JacobianTape(QuantumTape):
             # parameters. Simply return an empty Hessian.
             return np.zeros((len(params), len(params)), dtype=float)
 
+        # The parameter-shift Hessian implementation currently only supports
+        # the two-term parameter-shift rule. Raise an error for unsupported operations.
+        supported_ops = (
+            "RX",
+            "RY",
+            "RZ",
+            "Rot",
+            "PhaseShift",
+            "ControlledPhaseShift",
+            "MultiRZ",
+            "PauliRot",
+            "U1",
+            "U2",
+            "U3",
+            "SingleExcitationMinus",
+            "SingleExcitationPlus",
+            "DoubleExcitationMinus",
+            "DoubleExcitationPlus",
+        )
+
+        for idx, info in self._par_info.items():
+            op = info["op"]
+
+            if idx in self.trainable_params and op.name not in supported_ops:
+                raise ValueError(
+                    f"The operation {op.name} is currently not supported for the "
+                    f"parameter-shift Hessian.\nPlease decompose the operation in your "
+                    f"QNode by replacing it with '{op.__str__().replace('(', '.decomposition(')}'"
+                )
+
         # some gradient methods need the device or the device wires
         options["device"] = device
         options["dev_wires"] = device.wires


Test patch:
diff --git a/tests/interfaces/test_qnode_autograd.py b/tests/interfaces/test_qnode_autograd.py
index 720de5056..ee6485b98 100644
--- a/tests/interfaces/test_qnode_autograd.py
+++ b/tests/interfaces/test_qnode_autograd.py
@@ -827,6 +827,46 @@ class TestQNode:
         ]
         assert np.allclose(hess, expected_hess, atol=tol, rtol=0)
 
+    def test_hessian_unused_parameter(self, dev_name, diff_method, mocker, tol):
+        """Test hessian calculation of a scalar valued QNode"""
+        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.RY(x[0], wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        x = np.array([1.0, 2.0], requires_grad=True)
+        res = circuit(x)
+
+        a, b = x
+
+        expected_res = np.cos(a)
+        assert np.allclose(res, expected_res, atol=tol, rtol=0)
+
+        grad_fn = qml.grad(circuit)
+        g = grad_fn(x)
+
+        expected_g = [-np.sin(a), 0]
+        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(a), 0],
+            [0, 0],
+        ]
+        assert np.allclose(hess, expected_hess, atol=tol, rtol=0)
+
     def test_hessian_vector_valued(self, dev_name, diff_method, mocker, tol):
         """Test hessian calculation of a vector valued QNode"""
         if diff_method not in {"parameter-shift", "backprop"}:
diff --git a/tests/tape/test_qubit_param_shift.py b/tests/tape/test_qubit_param_shift.py
index 174d54710..6e3ac8d60 100644
--- a/tests/tape/test_qubit_param_shift.py
+++ b/tests/tape/test_qubit_param_shift.py
@@ -627,3 +627,65 @@ class TestHessian:
         hessian = tape.hessian(dev)
 
         assert hessian.shape == (0, 0)
+
+    @pytest.mark.parametrize("G", [qml.CRX, qml.CRY, qml.CRZ])
+    def test_controlled_rotation_error(self, G, tol):
+        """Test that attempting to perform the parameter-shift rule on the controlled rotation gates
+        results in an error."""
+        dev = qml.device("default.qubit", wires=2)
+        b = 0.123
+
+        with QubitParamShiftTape() as tape:
+            qml.QubitStateVector(np.array([1.0, -1.0]) / np.sqrt(2), wires=0)
+            G(b, wires=[0, 1])
+            qml.expval(qml.PauliX(0))
+
+        tape.trainable_params = {1}
+
+        res = tape.execute(dev)
+        assert np.allclose(res, -np.cos(b / 2), atol=tol, rtol=0)
+
+        with pytest.raises(ValueError, match="not supported for the parameter-shift Hessian"):
+            tape.hessian(dev, method="analytic")
+
+    @pytest.mark.parametrize("G", [qml.CRX, qml.CRY, qml.CRZ])
+    def test_controlled_rotation_second_derivative(self, G, tol):
+        """Test that the controlled rotation gates return the correct
+        second derivative if first decomposed."""
+        dev = qml.device("default.qubit", wires=2)
+        init_state = qml.numpy.array([1.0, -1.0], requires_grad=False) / np.sqrt(2)
+
+        @qml.qnode(dev)
+        def circuit(b):
+            qml.QubitStateVector(init_state, wires=0)
+            G(b, wires=[0, 1])
+            return qml.expval(qml.PauliX(0))
+
+        b = 0.123
+
+        res = circuit(b)
+        assert np.allclose(res, -np.cos(b / 2), atol=tol, rtol=0)
+
+        grad = qml.grad(qml.grad(circuit))(b)
+        expected = np.cos(b / 2) / 4
+        assert np.allclose(grad, expected, atol=tol, rtol=0)
+
+    def test_non_differentiable_controlled_rotation(self, tol):
+        """Tests that a non-differentiable controlled operation does not affect
+        the Hessian computation."""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        x = 0.6
+
+        with QubitParamShiftTape() as tape:
+            qml.RY(x, wires=0)
+            qml.CRY(np.pi / 2, wires=[0, 1])
+            qml.expval(qml.PauliX(0))
+
+        tape.trainable_params = {0}
+        hess = tape.hessian(dev)
+
+        expected_hess = np.array([-np.sin(x) / np.sqrt(2)])
+
+        assert np.allclose(hess, expected_hess, atol=tol, rtol=0)