Repository: PennyLaneAI/pennylane
Pull request #1110 (2021-02-24 20:33:31)

Add support for the parameter-shift hessian to the TF interface
This pull request ties in the parameter shift hessian calculation to the tensorflow interface.

Now you should be able to run:

```python
dev = qml.device('default.qubit', wires=1)
@qml.qnode(dev, interface='tf', diff_method='parameter-shift')
def circuit(x):
    qml.RX(x[0], wires=0)
    qml.RY(x[1], wires=0)
    return qml.expval(qml.PauliZ(0))

x = tf.Variable([0.1, 0.2], dtype=tf.float64)

with tf.GradientTape() as tape1:
    with tf.GradientTape() as tape2:
        y = circuit(x)
    grad = tape2.gradient(res, x)

tape1.jacobian(grad, x)
```

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [.github/CHANGELOG.md](/PennyLaneAI/pennylane/blob/master/.github/CHANGELOG.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1110?src=pr&el=h1) Report
> Merging [#1110](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1110?src=pr&el=desc) (7813490) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/c02ec7896dbf9feecd2baca1a303f2cb3f1cb3e1?el=desc) (c02ec78) will **increase** coverage by `3.14%`.
> The diff coverage is `96.96%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1110/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1110?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master    #1110      +/-   ##
==========================================
+ Coverage   94.81%   97.96%   +3.14%     
==========================================
  Files         142      142              
  Lines       10268    10288      +20     
==========================================
+ Hits         9736    10079     +343     
+ Misses        532      209     -323     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1110?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [pennylane/tape/jacobian\_tape.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1110/diff?src=pr&el=tree#diff-cGVubnlsYW5lL3RhcGUvamFjb2JpYW5fdGFwZS5weQ==) | `97.84% <83.33%> (-0.50%)` | :arrow_down: |
| [pennylane/interfaces/tf.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1110/diff?src=pr&el=tree#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvdGYucHk=) | `97.14% <100.00%> (+6.23%)` | :arrow_up: |
| [pennylane/tape/qubit\_param\_shift.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1110/diff?src=pr&el=tree#diff-cGVubnlsYW5lL3RhcGUvcXViaXRfcGFyYW1fc2hpZnQucHk=) | `100.00% <0.00%> (+0.82%)` | :arrow_up: |
| [pennylane/math/tensorbox.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1110/diff?src=pr&el=tree#diff-cGVubnlsYW5lL21hdGgvdGVuc29yYm94LnB5) | `95.97% <0.00%> (+1.34%)` | :arrow_up: |
| [pennylane/beta/devices/default\_tensor.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1110/diff?src=pr&el=tree#diff-cGVubnlsYW5lL2JldGEvZGV2aWNlcy9kZWZhdWx0X3RlbnNvci5weQ==) | `95.23% <0.00%> (+1.70%)` | :arrow_up: |
| [pennylane/devices/default\_qubit\_tf.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1110/diff?src=pr&el=tree#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdF90Zi5weQ==) | `90.66% <0.00%> (+2.66%)` | :arrow_up: |
| [pennylane/math/numpy\_box.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1110/diff?src=pr&el=tree#diff-cGVubnlsYW5lL21hdGgvbnVtcHlfYm94LnB5) | `98.55% <0.00%> (+2.89%)` | :arrow_up: |
| [pennylane/beta/devices/default\_tensor\_tf.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1110/diff?src=pr&el=tree#diff-cGVubnlsYW5lL2JldGEvZGV2aWNlcy9kZWZhdWx0X3RlbnNvcl90Zi5weQ==) | `87.50% <0.00%> (+3.12%)` | :arrow_up: |
| [pennylane/qnode.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1110/diff?src=pr&el=tree#diff-cGVubnlsYW5lL3Fub2RlLnB5) | `98.67% <0.00%> (+3.96%)` | :arrow_up: |
| [pennylane/interfaces/torch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1110/diff?src=pr&el=tree#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvdG9yY2gucHk=) | `100.00% <0.00%> (+4.68%)` | :arrow_up: |
| ... and [11 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1110/diff?src=pr&el=tree-more) | |

------

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


Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index a4eb07b99..83baf8109 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -2,6 +2,35 @@
 
 <h3>New features since last release</h3>
 
+- The TensorFlow interface now supports computing second derivatives and Hessians of hybrid quantum models.
+  Second derivatives are supported on both hardware and simulators.
+  [(#1110)](https://github.com/PennyLaneAI/pennylane/pull/1110) 
+
+  ```python
+  dev = qml.device('default.qubit', wires=1)
+  @qml.qnode(dev, interface='tf', diff_method='parameter-shift')
+  def circuit(x):
+      qml.RX(x[0], wires=0)
+      qml.RY(x[1], wires=0)
+      return qml.expval(qml.PauliZ(0))
+
+  x = tf.Variable([0.1, 0.2], dtype=tf.float64)
+
+  with tf.GradientTape() as tape1:
+      with tf.GradientTape() as tape2:
+          y = circuit(x)
+      grad = tape2.gradient(res, x)
+
+  hessian = tape1.jacobian(grad, x)
+  ```
+
+  To compute just the diagonal of the Hessian, the gradient of the
+  first derivatives can be taken:
+
+  ```python
+  hessian_diagonals = tape1.gradient(grad, x)
+  ```
+
 * Adds a new transform `qml.ctrl` that adds control wires to subroutines.
   [(#1157)](https://github.com/PennyLaneAI/pennylane/pull/1157)
 
diff --git a/pennylane/interfaces/tf.py b/pennylane/interfaces/tf.py
index 457f23296..286a7e8a7 100644
--- a/pennylane/interfaces/tf.py
+++ b/pennylane/interfaces/tf.py
@@ -138,34 +138,88 @@ class TFInterface(AnnotatedQueue):
         res = self.execute_device(args, input_kwargs["device"])
         self.set_parameters(all_params, trainable_only=False)
 
-        def grad(grad_output, **tfkwargs):
-            variables = tfkwargs.get("variables", None)
+        # The following dictionary caches the Jacobian and Hessian matrices,
+        # so that they can be re-used for different vjp/vhp computations
+        # within the same backpropagation call.
+        # This dictionary is tied to an instance of the inner function jacobian_product
+        # called within tf_tape.gradient or tf_tape.jacobian,
+        # via closure. Once tf_tape.gradient/ jacobian has returned, the jacobian_product instance
+        # will no longer be in scope and the memory will be freed.
+        saved_grad_matrices = {}
+
+        def _evaluate_grad_matrix(grad_matrix_fn):
+            """Convenience function for generating gradient matrices
+            for the given parameter values.
+
+            This function serves two purposes:
+
+            * Avoids duplicating logic surrounding parameter unwrapping/wrapping
+
+            * Takes advantage of closure, to cache computed gradient matrices via
+              the ``saved_grad_matrices`` dictionary, to avoid gradient matrices being
+              computed multiple redundant times.
+
+              This is particularly useful when differentiating vector-valued QNodes.
+              Because tensorflow requests the vector-grad matrix product,
+              and *not* the full grad matrix, differentiating vector-valued
+              functions will result in multiple backward passes.
+
+            Args:
+                grad_matrix_fn (str): Name of the gradient matrix function. Should correspond to an existing
+                    tape method. Currently allowed values include ``"jacobian"`` and ``"hessian"``.
+
+            Returns:
+                array[float]: the gradient matrix
+            """
+            if grad_matrix_fn in saved_grad_matrices:
+                return saved_grad_matrices[grad_matrix_fn]
 
             self.set_parameters(all_params_unwrapped, trainable_only=False)
-            jacobian = self.jacobian(input_kwargs["device"], params=args, **self.jacobian_options)
+            grad_matrix = getattr(self, grad_matrix_fn)(
+                input_kwargs["device"], params=args, **self.jacobian_options
+            )
             self.set_parameters(all_params, trainable_only=False)
 
-            jacobian = tf.constant(jacobian, dtype=self.dtype)
+            grad_matrix = tf.constant(grad_matrix, dtype=self.dtype)
+            saved_grad_matrices[grad_matrix_fn] = grad_matrix
+
+            return grad_matrix
+
+        def jacobian_product(dy, **tfkwargs):
+            variables = tfkwargs.get("variables", None)
+            dy_row = tf.reshape(dy, [1, -1])
+
+            @tf.custom_gradient
+            def jacobian(p):
+                def hessian_product(ddy, **tfkwargs):
+                    variables = tfkwargs.get("variables", None)
+                    hessian = _evaluate_grad_matrix("hessian")
+
+                    if self.output_dim == 1:
+                        hessian = tf.expand_dims(hessian, -1)
 
-            # Reshape gradient output array as a 2D row-vector.
-            grad_output_row = tf.reshape(grad_output, [1, -1])
+                    vhp = tf.cond(
+                        tf.rank(hessian) > 2,
+                        lambda: dy_row @ ddy @ hessian @ tf.transpose(dy_row),
+                        lambda: ddy @ hessian,
+                    )
 
-            # Calculate the vector-Jacobian matrix product, and unstack the output.
-            grad_input = tf.matmul(grad_output_row, jacobian)
-            grad_input = tf.unstack(tf.reshape(grad_input, [-1]))
+                    vhp = tf.unstack(tf.reshape(vhp, [-1]))
+                    return (vhp, variables) if variables is not None else vhp
 
-            if variables is not None:
-                return grad_input, variables
+                return _evaluate_grad_matrix("jacobian"), hessian_product
 
-            return grad_input
+            vjp = tf.matmul(dy_row, jacobian(params))
+            vjp = tf.unstack(tf.reshape(vjp, [-1]))
+            return (vjp, variables) if variables is not None else vjp
 
         if self.is_sampled:
-            return res, grad
+            return res, jacobian_product
 
         if res.dtype == np.dtype("object"):
             res = np.hstack(res)
 
-        return tf.convert_to_tensor(res, dtype=self.dtype), grad
+        return tf.convert_to_tensor(res, dtype=self.dtype), jacobian_product
 
     @classmethod
     def apply(cls, tape, dtype=tf.float64):
diff --git a/pennylane/tape/jacobian_tape.py b/pennylane/tape/jacobian_tape.py
index 2ecff9e56..6bb85f126 100644
--- a/pennylane/tape/jacobian_tape.py
+++ b/pennylane/tape/jacobian_tape.py
@@ -489,7 +489,7 @@ class JacobianTape(QuantumTape):
         >>> tape.jacobian(dev)
         array([], shape=(4, 0), dtype=float64)
         """
-        if any([m.return_type is State for m in self.measurements]):
+        if any(m.return_type is State for m in self.measurements):
             raise ValueError("The jacobian method does not support circuits that return the state")
 
         if self.is_sampled:
@@ -663,7 +663,7 @@ class JacobianTape(QuantumTape):
         >>> tape.hessian(dev)
         array([], shape=(0, 0), dtype=float64)
         """
-        if any([m.return_type is State for m in self.measurements]):
+        if any(m.return_type is State for m in self.measurements):
             raise ValueError("The Hessian method does not support circuits that return the state")
 
         method = options.get("method", "analytic")
@@ -730,11 +730,19 @@ class JacobianTape(QuantumTape):
 
             if hessian is None:
                 # create the Hessian matrix
-                hessian = np.zeros((len(params), len(params)), dtype=float)
+                if self.output_dim is not None:
+                    hessian = np.zeros(
+                        (len(params), len(params), np.prod(self.output_dim)), dtype=float
+                    )
+                else:
+                    hessian = np.zeros((len(params), len(params)), dtype=float)
 
             if i == j:
                 hessian[i, i] = g
             else:
                 hessian[i, j] = hessian[j, i] = g
 
+        if self.output_dim == 1:
+            hessian = np.squeeze(hessian, axis=-1)
+
         return hessian


Test patch:
diff --git a/tests/interfaces/test_qnode_tf.py b/tests/interfaces/test_qnode_tf.py
index 475a30a3a..45bcca3a2 100644
--- a/tests/interfaces/test_qnode_tf.py
+++ b/tests/interfaces/test_qnode_tf.py
@@ -32,7 +32,7 @@ from pennylane.tape import JacobianTape
     ],
 )
 class TestQNode:
-    """Same tests as above, but this time via the QNode interface!"""
+    """Tests the tensorflow interface used with a QNode."""
 
     def test_execution_no_interface(self, dev_name, diff_method):
         """Test execution works without an interface, and that trainable parameters
@@ -525,6 +525,208 @@ class TestQNode:
         assert res.shape == (2, 10)
         assert isinstance(res, tf.Tensor)
 
+    def test_second_derivative(self, dev_name, diff_method, mocker, tol):
+        """Test second derivative 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="tf")
+        def circuit(x):
+            qml.RY(x[0], wires=0)
+            qml.RX(x[1], wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        x = tf.Variable([1.0, 2.0], dtype=tf.float64)
+
+        with tf.GradientTape() as tape1:
+            with tf.GradientTape() as tape2:
+                res = circuit(x)
+            g = tape2.gradient(res, x)
+            res2 = tf.reduce_sum(g)
+
+        spy = mocker.spy(JacobianTape, "hessian")
+        g2 = tape1.gradient(res2, x)
+
+        if diff_method == "parameter-shift":
+            spy.assert_called_once()
+        elif diff_method == "backprop":
+            spy.assert_not_called()
+
+        a, b = x * 1.0
+
+        expected_res = tf.cos(a) * tf.cos(b)
+        assert np.allclose(res, expected_res, atol=tol, rtol=0)
+
+        expected_g = [-tf.sin(a) * tf.cos(b), -tf.cos(a) * tf.sin(b)]
+        assert np.allclose(g, expected_g, atol=tol, rtol=0)
+
+        expected_g2 = [-tf.cos(a) * tf.cos(b) + tf.sin(a) * tf.sin(b), tf.sin(a) * tf.sin(b) - tf.cos(a) * tf.cos(b)]
+        assert np.allclose(g2, expected_g2, atol=tol, rtol=0)
+
+    def test_hessian(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="tf")
+        def circuit(x):
+            qml.RY(x[0], wires=0)
+            qml.RX(x[1], wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        x = tf.Variable([1.0, 2.0], dtype=tf.float64)
+
+        with tf.GradientTape() as tape1:
+            with tf.GradientTape() as tape2:
+                res = circuit(x)
+            g = tape2.gradient(res, x)
+
+        spy = mocker.spy(JacobianTape, "hessian")
+        hess = tape1.jacobian(g, x)
+
+        if diff_method == "parameter-shift":
+            spy.assert_called_once()
+        elif diff_method == "backprop":
+            spy.assert_not_called()
+
+        a, b = x * 1.0
+
+        expected_res = tf.cos(a) * tf.cos(b)
+        assert np.allclose(res, expected_res, atol=tol, rtol=0)
+
+        expected_g = [-tf.sin(a) * tf.cos(b), -tf.cos(a) * tf.sin(b)]
+        assert np.allclose(g, expected_g, atol=tol, rtol=0)
+
+        expected_hess = [
+            [-tf.cos(a) * tf.cos(b), tf.sin(a) * tf.sin(b)],
+            [tf.sin(a) * tf.sin(b), -tf.cos(a) * tf.cos(b)]
+        ]
+        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"}:
+            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.RY(x[0], wires=0)
+            qml.RX(x[1], wires=0)
+            return qml.probs(wires=0)
+
+        x = tf.Variable([1.0, 2.0], dtype=tf.float64)
+
+        with tf.GradientTape(persistent=True) as tape1:
+            with tf.GradientTape(persistent=True) as tape2:
+                res = circuit(x)
+
+            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 = [
+            0.5 + 0.5 * tf.cos(a) * tf.cos(b),
+            0.5 - 0.5 * tf.cos(a) * tf.cos(b)
+        ]
+        assert np.allclose(res, expected_res, atol=tol, rtol=0)
+
+        expected_g = [
+            [-0.5 * tf.sin(a) * tf.cos(b), -0.5 * tf.cos(a) * tf.sin(b)],
+            [0.5 * tf.sin(a) * tf.cos(b), 0.5 * tf.cos(a) * tf.sin(b)]
+        ]
+        assert np.allclose(g, expected_g, atol=tol, rtol=0)
+
+        expected_hess = [
+            [
+                [-0.5 * tf.cos(a) * tf.cos(b), 0.5 * tf.sin(a) * tf.sin(b)],
+                [0.5 * tf.sin(a) * tf.sin(b), -0.5 * tf.cos(a) * tf.cos(b)]
+            ],
+            [
+                [0.5 * tf.cos(a) * tf.cos(b), -0.5 * tf.sin(a) * tf.sin(b)],
+                [-0.5 * tf.sin(a) * tf.sin(b), 0.5 * tf.cos(a) * tf.cos(b)]
+            ]
+        ]
+
+        np.testing.assert_allclose(hess, expected_hess, atol=tol, rtol=0, verbose=True)
+
+    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"}:
+            pytest.skip("Test only supports parameter-shift or backprop")
+
+        dev = qml.device(dev_name, wires=2)
+
+        @qnode(dev, diff_method=diff_method, interface="tf")
+        def circuit(x):
+            qml.RY(x[0], wires=0)
+            qml.RX(x[1], wires=0)
+            qml.RY(x[0], wires=1)
+            qml.RX(x[1], wires=1)
+            return qml.expval(qml.PauliZ(0)), qml.probs(wires=1)
+
+        x = tf.Variable([1.0, 2.0], dtype=tf.float64)
+        res = circuit(x)
+
+        with tf.GradientTape(persistent=True) as tape1:
+            with tf.GradientTape(persistent=True) as tape2:
+                res = circuit(x)
+
+            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 = [
+            tf.cos(a) * tf.cos(b),
+            0.5 + 0.5 * tf.cos(a) * tf.cos(b),
+            0.5 - 0.5 * tf.cos(a) * tf.cos(b)
+        ]
+        assert np.allclose(res, expected_res, atol=tol, rtol=0)
+
+        expected_g = [
+            [-tf.sin(a) * tf.cos(b), -tf.cos(a) * tf.sin(b)],
+            [-0.5 * tf.sin(a) * tf.cos(b), -0.5 * tf.cos(a) * tf.sin(b)],
+            [0.5 * tf.sin(a) * tf.cos(b), 0.5 * tf.cos(a) * tf.sin(b)]
+        ]
+        assert np.allclose(g, expected_g, atol=tol, rtol=0)
+
+        expected_hess = [
+            [
+                [-tf.cos(a) * tf.cos(b), tf.sin(a) * tf.sin(b)],
+                [tf.sin(a) * tf.sin(b), -tf.cos(a) * tf.cos(b)]
+            ],
+            [
+                [-0.5 * tf.cos(a) * tf.cos(b), 0.5 * tf.sin(a) * tf.sin(b)],
+                [0.5 * tf.sin(a) * tf.sin(b), -0.5 * tf.cos(a) * tf.cos(b)]
+            ],
+            [
+                [0.5 * tf.cos(a) * tf.cos(b), -0.5 * tf.sin(a) * tf.sin(b)],
+                [-0.5 * tf.sin(a) * tf.sin(b), 0.5 * tf.cos(a) * tf.cos(b)]
+            ]
+        ]
+        np.testing.assert_allclose(hess, expected_hess, atol=tol, rtol=0, verbose=True)
+
 
 def qtransform(qnode, a, framework=tf):
     """Transforms every RY(y) gate in a circuit to RX(-a*cos(y))"""
diff --git a/tests/tape/test_qubit_param_shift.py b/tests/tape/test_qubit_param_shift.py
index 045c65389..89097ee9d 100644
--- a/tests/tape/test_qubit_param_shift.py
+++ b/tests/tape/test_qubit_param_shift.py
@@ -579,6 +579,33 @@ class TestHessian:
 
         assert np.allclose(autograd_val, manualgrad_val, atol=tol, rtol=0)
 
+    def test_vector_output(self, tol):
+        """Tests that a vector valued output tape has a hessian with the proper result. """
+
+        dev = qml.device('default.qubit', wires=1)
+
+        x = np.array([1.0, 2.0])
+
+        with QubitParamShiftTape() as tape:
+            qml.RY(x[0], wires=0)
+            qml.RX(x[1], wires=0)
+            qml.probs(wires=[0])
+
+        hess = tape.hessian(dev)
+
+        expected_hess = expected_hess = np.array([
+            [
+                [-0.5 * np.cos(x[0]) * np.cos(x[1]), 0.5 * np.cos(x[0]) * np.cos(x[1])],
+                [ 0.5 * np.sin(x[0]) * np.sin(x[1]), -0.5 * np.sin(x[0]) * np.sin(x[1])]
+            ],
+            [
+                [0.5 * np.sin(x[0]) * np.sin(x[1]), -0.5 * np.sin(x[0]) * np.sin(x[1])],
+                [-0.5 * np.cos(x[0]) * np.cos(x[1]), 0.5 * np.cos(x[0]) * np.cos(x[1])]
+            ]
+        ])
+
+        assert np.allclose(hess, expected_hess, atol=tol, rtol=0)
+
     def test_no_trainable_params_hessian(self):
         """Test that an empty Hessian is returned when there are no trainable
         parameters."""