Repository: PennyLaneAI/pennylane (python)

Issue:
Add a parameter shift Hessian transform
### Feature details

Currently, it is possible to compute the parameter-shift Hessian of a QNode by using an autodifferentiation framework, and asking for the Jacobian of the gradient (or the Jacobian of the Jacobian, for a vector-valued QNode):

```python
dev = qml.device("default.qubit", wires=2)

@qml.beta.qnode(dev, diff_method=qml.gradients.param_shift, max_diff=2)
def circuit(x):
    qml.RX(x[0], wires=0)
    qml.RY(x[1], wires=0)
    return qml.probs(wires=0)
```
```pycon
>>> x = np.array([0.1, 0.2], requires_grad=True)
>>> hess_fn = qml.jacobian(qml.jacobian(circuit))
>>> hess_fn(x)
[[[-0.48758516  0.00991692]
  [ 0.00991692 -0.48758516]]
 [[ 0.48758516 -0.00991692]
  [-0.00991692  0.48758516]]]
```

This works well if integrating the QNode into an autodifferentiation pipeline.

However, if you instead wish to compute the Hessian of a QNode in isolation, this is not ideal:

- Autodifferentiation frameworks, as they are designed for backpropagation, always evaluate the QNode, followed by the gradient, *before* computing the second derivatives. As a result, if you are *only interested in the Hessian*, there will be redundant evaluations.

- Further, if we know we only require the second derivatives, there are tricks we can use to further minimize the number of quantum evaluations. See https://arxiv.org/abs/2008.06517 and https://arxiv.org/abs/2107.12390.

- Finally, having a separate transform for computing the Hessian enables better analysis, e.g., the ability to check view the generated 'Hessian tapes', or to query the number of evaluations.

As a result, we should add a [gradient transform](https://pennylane.readthedocs.io/en/latest/code/qml_gradients.html),

```python
>>> qml.gradients.param_shift_hessian(qnode)(*args, **kwargs)
```

that returns the Hessian of a QNode _directly_, using one or more of the approaches given in https://arxiv.org/abs/2008.06517 and https://arxiv.org/abs/2107.12390.

### Implementation

- A new file should be added, `pennylane/gradients/param_shift_hessian.py`. The function should be written to be *autodiff framework agnostic* using `qml.math`, so that it will work on QNodes defined for all interfaces.

- The new function should be included in the `pennylane/gradients/__init__.py` module docstring

- Test should be added in the file `tests/gradients/test_param_shift_hessian.py`. Things to check work correctly:

  - Operators with 2-term shift rules (`qml.RX()`) should return the correct second derivatives
  - Operators with 4-term shift rules (`qml.CRX()`) should return the correct second derivatives
  - The transform should generate the minimal number of quantum evaluations.
  - The transform should be *differentiable*. That is, you should be able to do
  
    ```python
    >>> qml.grad(qml.gradients.param_shift_hessian(qnode))(weights)
    ```
    to get 3rd derivatives. Furthermore, this differentiability should work in all interfaces, including TF and Torch.



### How important would you say this feature is?

2: Somewhat important. Needed this quarter.

### Additional information

_No response_

Pull request description:
Add a parameter-shift Hessian transformation
**Context:** Currently, it is possible to compute the parameter-shift Hessian of a QNode by using an autodifferentiation framework, and asking for the Jacobian of the gradient (or the Jacobian of the Jacobian, for a vector-valued QNode). However, this leads to redundant evaluations from computing the QNode and first-derivative of the QNode.

**Description of the Change:** Adds a new transformation that directly computes the second-derivative (Hessian) of QNodes via parameter-shift rules.

**Benefits:** Fewer quantum evaluations, can directly inspect the "Hessian tapes".

**Possible Drawbacks:**

**Related GitHub Issues:** Closes #1752


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index f8d799a09..94016a23a 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -214,6 +214,39 @@
   The adjoint method requires memory for 4 independent state vectors, which corresponds roughly
   to storing a state vector of a system with 2 additional qubits.
 
+* A new method `qml.gradients.param_shift_hessian` has been added to directly compute the Hessian
+  (2nd order partial derivative matrix) of QNodes and QuantumTapes. The method generates
+  parameter-shifted tapes which allow the Hessian to be computed analytically on hardware and
+  software devices. Compared to using an auto-differentiation framework to compute the Hessian
+  via parameter shifts, this method will use fewer device invocations and can be used to inspect
+  the parameter-shifted "Hessian tapes" directly. The method remains fully differentiable on all
+  supported PennyLane interfaces.
+
+  Additionally, the parameter-shift Hessian comes with a new batch transform decorator
+  `@qml.gradients.hessian_transform`, which can be used to create custom Hessian methods.
+  [(#1884)](https://github.com/PennyLaneAI/pennylane/pull/1884)
+
+  The following code demonstrates how to use the parameter-shift Hessian:
+
+  ```python
+  dev = qml.device("default.qubit", wires=2)
+
+  @qml.qnode(dev)
+  def circuit(x):
+      qml.RX(x[0], wires=0)
+      qml.RY(x[1], wires=0)
+      return qml.expval(qml.PauliZ(0))
+
+  x = np.array([0.1, 0.2], requires_grad=True)
+
+  hessian = qml.gradients.param_shift_hessian(circuit)(x)
+  ```
+  ```pycon
+  >>> hessian
+  tensor([[-0.97517033,  0.01983384],
+          [ 0.01983384, -0.97517033]], requires_grad=True)
+  ```
+
 <h3>Improvements</h3>
 
 * The function `qml.math.safe_squeeze` is introduced and `gradient_transform` allows
diff --git a/pennylane/gradients/__init__.py b/pennylane/gradients/__init__.py
index 79142392c..38a0d4581 100644
--- a/pennylane/gradients/__init__.py
+++ b/pennylane/gradients/__init__.py
@@ -45,6 +45,7 @@ Gradient transforms
     finite_diff
     param_shift
     param_shift_cv
+    param_shift_hessian
 
 Custom gradients
 ^^^^^^^^^^^^^^^^
@@ -53,6 +54,7 @@ Custom gradients
     :toctree: api
 
     gradient_transform
+    hessian_transform
 
 Utility functions
 ^^^^^^^^^^^^^^^^^
@@ -283,11 +285,14 @@ import pennylane as qml
 from . import finite_difference
 from . import parameter_shift
 from . import parameter_shift_cv
+from . import param_shift_hessian
 
 from .gradient_transform import gradient_transform
+from .hessian_transform import hessian_transform
 from .finite_difference import finite_diff, finite_diff_coeffs, generate_shifted_tapes
 from .parameter_shift import param_shift
 from .parameter_shift_cv import param_shift_cv
+from .param_shift_hessian import param_shift_hessian
 from .vjp import compute_vjp, batch_vjp, vjp
 from .hamiltonian_grad import hamiltonian_grad
 from .general_shift_rules import (
diff --git a/pennylane/gradients/hessian_transform.py b/pennylane/gradients/hessian_transform.py
new file mode 100644
index 000000000..041dd1847
--- /dev/null
+++ b/pennylane/gradients/hessian_transform.py
@@ -0,0 +1,158 @@
+# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""This module contains utilities for defining custom Hessian transforms,
+including a decorator for specifying Hessian expansions."""
+import warnings
+
+import pennylane as qml
+
+from pennylane.transforms.tape_expand import expand_invalid_trainable
+
+
+class hessian_transform(qml.batch_transform):
+    """Decorator for defining quantum Hessian transforms.
+
+    Quantum Hessian transforms are a specific case of :class:`~.batch_transform`s,
+    similar to the :class:`~.gradient_transform`. Hessian transforms compute the
+    second derivative of a quantum function.
+    All quantum Hessian transforms accept a tape, and output a batch of tapes to
+    be independently executed on a quantum device, alongside a post-processing
+    function to return the result.
+
+    Args:
+        expand_fn (function): An expansion function (if required) to be applied to the
+            input tape before the Hessian computation takes place. If not provided,
+            the default expansion function simply expands all operations that
+            have ``Operation.grad_method=None`` until all resulting operations
+            have a defined gradient method.
+        differentiable (bool): Specifies whether the Hessian transform is differentiable
+            or not. A transform may be non-differentiable if it does not use an autodiff
+            framework for its tensor manipulations. In such a case, setting
+            ``differentiable=False`` instructs the decorator to mark the output as
+            'constant', reducing potential overhead.
+        hybrid (bool): Specifies whether classical processing inside a QNode
+            should be taken into account when transforming a QNode.
+
+            - If ``True``, and classical processing is detected, the Jacobian of the
+              classical processing will be computed and included. When evaluated, the
+              returned Hessian will be with respect to the QNode arguments.
+
+            - If ``False``, any internal QNode classical processing will be **ignored**.
+              When evaluated, the returned Hessian will be with respect to the **gate**
+              arguments, and not the QNode arguments.
+
+    Supported Hessian transforms must be of the following form:
+
+    .. code-block:: python
+
+        @hessian_transform
+        def my_custom_hessian(tape, **kwargs):
+            ...
+            return hessian_tapes, processing_fn
+
+    where:
+
+    - ``tape`` (*QuantumTape*): the input quantum tape to compute the Hessian of
+
+    - ``hessian_tapes`` (*List[QuantumTape]*): is a list of output tapes to be
+      evaluated. If this list is empty, no quantum evaluations will be made.
+
+    - ``processing_fn`` is a processing function to be applied to the output of the
+      evaluated ``hessian_tapes``. It should accept a list of numeric results with
+      length ``len(hessian_tapes)``, and return the Hessian matrix.
+
+    Once defined, the quantum Hessian transform can be used as follows:
+
+    >>> hessian_tapes, processing_fn = my_custom_hessian(tape, *hessian_kwargs)
+    >>> res = execute(tapes, dev, interface="autograd", gradient_fn=qml.gradients.param_shift)
+    >>> jacobian = processing_fn(res)
+
+    Alternatively, Hessian transforms can be applied directly to QNodes, in which case
+    the execution is implicit:
+
+    >>> fn = my_custom_hessian(qnode, *hessian_kwargs)
+    >>> fn(weights)  # transformed function takes the same arguments as the QNode
+    1.2629730888100839
+
+    .. note::
+
+        The input tape might have parameters of various types, including NumPy arrays,
+        JAX DeviceArrays, and TensorFlow and PyTorch tensors.
+
+        If the Hessian transform is written in a autodiff-compatible manner, either by
+        using a framework such as Autograd or TensorFlow, or by using ``qml.math`` for
+        tensor manipulation, then higher-order derivatives will also be supported.
+
+        Alternatively, you may use the ``tape.unwrap()`` context manager to temporarily
+        convert all tape parameters to NumPy arrays and floats:
+
+        >>> with tape.unwrap():
+        ...     params = tape.get_parameters()  # list of floats
+    """
+
+    def __init__(
+        self, transform_fn, expand_fn=expand_invalid_trainable, differentiable=True, hybrid=True
+    ):
+        self.hybrid = hybrid
+        super().__init__(transform_fn, expand_fn=expand_fn, differentiable=differentiable)
+
+    def default_qnode_wrapper(self, qnode, targs, tkwargs):
+        # Here, we overwrite the QNode execution wrapper in order to take into account
+        # that classical processing may be present inside the QNode.
+        hybrid = tkwargs.pop("hybrid", self.hybrid)
+        _wrapper = super().default_qnode_wrapper(qnode, targs, tkwargs)
+        cjac_fn = qml.transforms.classical_jacobian(qnode, expand_fn=self.expand_fn)
+
+        def hessian_wrapper(*args, **kwargs):
+            qhess = _wrapper(*args, **kwargs)
+
+            if any(m.return_type is qml.operation.Probability for m in qnode.qtape.measurements):
+                qhess = qml.math.squeeze(qhess)
+
+            if not hybrid:
+                return qhess
+
+            kwargs.pop("shots", False)
+            cjac = cjac_fn(*args, **kwargs)
+
+            has_single_arg = False
+            if not isinstance(cjac, tuple):
+                has_single_arg = True
+                cjac = (cjac,)
+
+            # The classical Jacobian for each argument has shape:
+            #   (# gate_args, *qnode_arg_shape)
+            # The Jacobian needs to be contracted twice with the quantum Hessian of shape:
+            #   (*qnode_output_shape, # gate_args, # gate_args)
+            # The result should then have the shape:
+            #   (*qnode_output_shape, *qnode_arg_shape, *qnode_arg_shape)
+            hessians = []
+            for jac in cjac:
+                if jac is not None:
+                    # Check for a Jacobian equal to the identity matrix.
+                    shape = qml.math.shape(jac)
+                    is_square = shape == (1,) or (len(shape) == 2 and shape[0] == shape[1])
+                    if is_square and qml.math.allclose(jac, qml.numpy.eye(shape[0])):
+                        hessians.append(qhess)
+                        continue
+
+                    num_arg_dims = len(qml.math.shape(jac)) - 1  # number of axes in qnode_arg_shape
+                    hess = qml.math.tensordot(qhess, jac, [[-1], [0]])
+                    hess = qml.math.tensordot(hess, jac, [[-1 - num_arg_dims], [0]])
+
+                    hessians.append(hess)
+
+            return hessians[0] if has_single_arg else tuple(hessians)
+
+        return hessian_wrapper
diff --git a/pennylane/gradients/param_shift_hessian.py b/pennylane/gradients/param_shift_hessian.py
new file mode 100644
index 000000000..f3c8e39a5
--- /dev/null
+++ b/pennylane/gradients/param_shift_hessian.py
@@ -0,0 +1,337 @@
+# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+This module contains functions for computing the parameter-shift hessian
+of a qubit-based quantum tape.
+"""
+import numpy as np
+
+import pennylane as qml
+
+from .parameter_shift import _gradient_analysis
+from .hessian_transform import hessian_transform
+
+
+def _process_hessian_recipe(hessian_recipe, tol=1e-10):
+    """Utility function to process Hessian recipes."""
+
+    hessian_recipe = np.array(hessian_recipe).T
+    # remove all small coefficients, shifts, and multipliers
+    hessian_recipe[np.abs(hessian_recipe) < tol] = 0
+    # remove columns where the coefficients are 0
+    hessian_recipe = hessian_recipe[:, :, ~(hessian_recipe[0, 0] == 0)]
+    # sort columns according to abs(shift2) then abs(shift1)
+    hessian_recipe = hessian_recipe[:, :, np.lexsort(np.abs(hessian_recipe)[:, -1])]
+
+    return hessian_recipe[0, 0], hessian_recipe[0, 1], hessian_recipe[:, 2].T
+
+
+def generate_multishifted_tapes(tape, idx, shifts):
+    r"""Generate a list of tapes where the corresponding trainable parameter
+    indices have been shifted by the values given.
+
+    Args:
+        tape (.QuantumTape): input quantum tape
+        idx (list[int]): trainable parameter indices to shift the parameters of
+        shifts (list[list[float or int]]): nested list of shift values, each
+            list containing a value for each index
+
+    Returns:
+        list[QuantumTape]: List of quantum tapes. Each tape has multiple parameters
+            (indicated by ``idx``) shifted by the values of ``shifts``. The length
+            of the returned list of tapes will match the length of ``shifts``.
+    """
+    params = list(tape.get_parameters())
+    tapes = []
+
+    for shift in shifts:
+        new_params = params.copy()
+        shifted_tape = tape.copy(copy_operations=True)
+
+        for id_, s in zip(idx, shift):
+            new_params[id_] = new_params[id_] + qml.math.convert_like(s, new_params[id_])
+
+        shifted_tape.set_parameters(new_params)
+        tapes.append(shifted_tape)
+
+    return tapes
+
+
+def compute_hessian_tapes(tape, diff_methods, f0=None):
+    r"""Generate the Hessian tapes that are used in the computation of the second derivative of a
+    quantum tape, using analytical parameter-shift rules to do so exactly. Also define a
+    post-processing function to combine the results of evaluating the Hessian tapes.
+
+    Args:
+        tape (.QuantumTape): input quantum tape
+        diff_methods (list[string]): The differentiation method to use for each trainable parameter.
+            Can be "A" or "0", where "A" is the analytical parameter shift rule and "0" indicates
+            a 0 derivative (that is the parameter does not affect the tape's output).
+        f0 (tensor_like[float] or None): Output of the evaluated input tape. If provided,
+            and the Hessian tapes include the original input tape, the 'f0' value is used
+            instead of evaluating the input tape, reducing the number of device invocations.
+
+    Returns:
+        tuple[list[QuantumTape], function]: A tuple containing a list of generated tapes, in
+            addition to a post-processing function to be applied to the results of the evaluated
+            tapes.
+    """
+    h_dim = tape.num_params
+
+    hessian_tapes = []
+    hessian_coeffs = []
+    unshifted_coeffs = {}
+    shapes = []
+
+    # The Hessian for a 2-term parameter-shift rule can be expressed via the following recipes.
+    # Off-diagonal elements of the Hessian require shifts to two different parameter indices.
+    # A recipe can thus be expressed via the tape patterns:
+    #       [[coeff, dummy], [mult, dummy], [shift1, shift2]]    (dummy values for ndarray creation)
+    # Each corresponding to one term in the parameter-shift formula:
+    #       didj f(x) = coeff * f(mult*x + shift1*ei + shift2*ej) + ...
+    diag_recipe = [[[0.5], [1], [np.pi]], [[-0.5], [1], [0]]]
+    off_diag_recipe = [
+        [[0.25, 1], [1, 1], [np.pi / 2, np.pi / 2]],
+        [[-0.25, 1], [1, 1], [-np.pi / 2, np.pi / 2]],
+        [[-0.25, 1], [1, 1], [np.pi / 2, -np.pi / 2]],
+        [[0.25, 1], [1, 1], [-np.pi / 2, -np.pi / 2]],
+    ]
+
+    for i in range(h_dim):
+        # optimization: skip partial derivatives that are zero
+        if diff_methods[i] == "0":
+            # We may fill in coefficients and shapes for the entire row.
+            hessian_coeffs.extend([[]] * h_dim)
+            shapes.extend(((i, j), 0) for j in range(h_dim))
+            continue
+
+        # Prefill Hessian in the lower triangular part for optimization below
+        hessian_coeffs.extend([[]] * i)
+        shapes.extend(((i, j), 0) for j in range(i))
+
+        # optimization: only generate tapes for upper triangular matrix (j >= i)
+        for j in range(i, h_dim):
+            # optimization: skip partial derivatives that are zero
+            if diff_methods[j] == "0":
+                hessian_coeffs.append([])
+                shapes.append(((i, j), 0))
+                continue
+
+            recipe = diag_recipe if i == j else off_diag_recipe
+            coeffs, _, shifts = _process_hessian_recipe(recipe)
+
+            # optimization: only compute the unshifted tape once
+            if all(np.array(shifts[0]) == 0):
+                if not unshifted_coeffs and f0 is None:
+                    hessian_tapes.insert(0, tape)
+
+                unshifted_coeffs[(i, j)] = coeffs[0]
+                coeffs, shifts = coeffs[1:], shifts[1:]
+
+            # generate the Hessian tapes
+            hessian_coeffs.append(coeffs)
+            h_tapes = generate_multishifted_tapes(tape, (i, j), shifts)
+
+            hessian_tapes.extend(h_tapes)
+            shapes.append(((i, j), len(h_tapes)))
+
+    def processing_fn(results):
+        # The first results dimension is the number of terms/tapes in the parameter-shift
+        # rule, the remaining ones are the QNode output dimensions.
+        out_dim = qml.math.shape(qml.math.stack(results))[1:]
+        # The desired shape of the Hessian is:
+        #       (QNode output dimensions, # trainable gate args, # trainable gate args),
+        # but first we accumulate all elements into a list, since no array assingment is possible.
+        hessian = []
+        # Keep track of tape results already consumed.
+        start = 1 if unshifted_coeffs and f0 is None else 0
+        # Results of the unshifted tape.
+        r0 = results[0] if start == 1 else f0
+
+        for k, ((i, j), s) in enumerate(shapes):
+            res = results[start : start + s]
+            start = start + s
+
+            # Compute the elements of the Hessian as the linear combination of
+            # results and coefficients, barring optimization cases.
+            if j < i:
+                g = hessian[j * h_dim + i]
+            elif s == 0:
+                g = qml.math.zeros(out_dim)
+            else:
+                res = qml.math.stack(res)
+                g = qml.math.tensordot(
+                    res, qml.math.convert_like(hessian_coeffs[k], res), [[0], [0]]
+                )
+                if (i, j) in unshifted_coeffs:
+                    g += unshifted_coeffs[(i, j)] * r0
+
+            hessian.append(g)
+
+        # Reshape the Hessian to have the dimensions of the QNode output on the outside, that is:
+        #         (h_dim, h_dim, out_dim) -> (out_dim, h_dim, h_dim)
+        hessian = qml.math.reshape(qml.math.stack(hessian), (h_dim, h_dim) + out_dim)
+        reordered_axes = list(range(2, len(out_dim) + 2)) + [0, 1]
+        hessian = qml.math.transpose(hessian, axes=reordered_axes)
+
+        return qml.math.squeeze(hessian)
+
+    return hessian_tapes, processing_fn
+
+
+@hessian_transform
+def param_shift_hessian(tape, f0=None):
+    r"""Transform a QNode to compute the parameter-shift Hessian with respect to its trainable
+    parameters.
+
+    Use this transform to explicitly generate and explore parameter-shift circuits for computing
+    the Hessian of QNodes directly, without computing first derivatives.
+
+    For second-order derivatives of more complicated cost functions, please consider using your
+    chosen autodifferentiation framework directly, by chaining gradient computations:
+
+    >>> qml.jacobian(qml.grad(cost))(weights)
+
+    .. note::
+
+        Currently, parametric gates are only supported if they obey a two-term shift rule,
+        which includes the following operations:
+
+        "RX", "RY", "RZ", "Rot", "PhaseShift", "ControlledPhaseShift", "MultiRZ", "PauliRot",
+        "U1", "U2", "U3", "SingleExcitationMinus", "SingleExcitationPlus", "DoubleExcitationMinus",
+        "DoubleExcitationPlus", "OrbitalRotation".
+
+    Args:
+        tape (pennylane.QNode or .QuantumTape): quantum tape or QNode to differentiate
+        f0 (tensor_like[float] or None): Output of the evaluated input tape. If provided,
+            and the Hessian tapes include the original input tape, the 'f0' value is used
+            instead of evaluating the input tape, reducing the number of device invocations.
+
+    Returns:
+        tensor_like or tuple[tensor_like] or tuple[list[QuantumTape], function]:
+
+        - If the input is a QNode with a single trainable argument, a tensor representing the
+          Hessian of size ``(*QNode output dimensions, *QNode input dimensions, *QNode input dimensions)``
+          is returned.
+
+        - If the input is a QNode with multiple trainable arguments, a tuple of Hessian tensors is
+          returned, one for each argument.
+
+        - If the input is a tape, a tuple containing the list of parameter-shifted tapes, and a
+          post-processing function to be applied to the evaluated tapes, is returned.
+
+        Note: By default a QNode with the keyword ``hybrid=True`` computes derivates with respect to
+        QNode arguments, which can include classical computations on those arguments before they are
+        passed to quantum operations. The "purely quantum" Hessian can instead be obtained with
+        ``hybrid=False``, which is then computed with respect to the gate arguments and produces a
+        result of shape ``(*QNode output dimensions, # gate arguments, # gate arguments)``.
+
+    **Example**
+
+    Applying the Hessian transform to a QNode computes its Hessian tensor:
+
+    >>> dev = qml.device("default.qubit", wires=2)
+    >>> @qml.qnode(dev)
+    ... def circuit(x):
+    ...     qml.RX(x[0], wires=0)
+    ...     qml.RY(x[1], wires=0)
+    ...     return qml.expval(qml.PauliZ(0))
+    >>> x = np.array([0.1, 0.2], requires_grad=True)
+    >>> qml.gradients.param_shift_hessian(circuit)(x)
+    tensor([[-0.97517033,  0.01983384],
+            [ 0.01983384, -0.97517033]], requires_grad=True)
+
+    .. UsageDetails::
+
+        The Hessian transform can also be applied to a quantum tape, instead producing the
+        parameter-shifted Hessian tapes and a post-processing function to combine execution
+        results:
+
+        >>> circuit(x)  # generate the QuantumTape inside the QNode
+        >>> tape = circuit.qtape
+        >>> hessian_tapes, postproc_fn = qml.gradients.param_shift_hessian(tape)
+        >>> hessian_tapes
+        [<JacobianTape: wires=[0], params=2>,
+            <JacobianTape: wires=[0], params=2>,
+            <JacobianTape: wires=[0], params=2>,
+            <JacobianTape: wires=[0], params=2>,
+            <JacobianTape: wires=[0], params=2>,
+            <JacobianTape: wires=[0], params=2>,
+            <JacobianTape: wires=[0], params=2>]
+        >>> postproc_fn(qml.execute(hessian_tapes, dev, None))
+        array([[-0.97517033,  0.01983384],
+                [ 0.01983384, -0.97517033]])
+
+        The Hessian tapes can be inspected via their draw function, which reveals the different
+        gate arguments generated from parameter-shift rules:
+
+        >>> for h_tape in hessian_tapes:
+        ...     print(h_tape.draw())
+        0: ──RX(0.1)──RY(0.2)──┤ ⟨Z⟩
+        0: ──RX(3.24)──RY(0.2)──┤ ⟨Z⟩
+        0: ──RX(1.67)──RY(1.77)──┤ ⟨Z⟩
+        0: ──RX(-1.47)──RY(1.77)──┤ ⟨Z⟩
+        0: ──RX(1.67)──RY(-1.37)──┤ ⟨Z⟩
+        0: ──RX(-1.47)──RY(-1.37)──┤ ⟨Z⟩
+        0: ──RX(0.1)──RY(3.34)──┤ ⟨Z⟩
+    """
+
+    # Perform input validation before generating tapes.
+    if any(m.return_type is qml.operation.State for m in tape.measurements):
+        raise ValueError(
+            "Computing the Hessian of circuits that return the state is not supported."
+        )
+
+    # The parameter-shift Hessian implementation currently doesn't support variance measurements.
+    if any(m.return_type is qml.operation.Variance for m in tape.measurements):
+        raise ValueError(
+            "Computing the Hessian of circuits that return variances is currently not supported."
+        )
+
+    if not tape.trainable_params:
+        return [], lambda _: []
+
+    # 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",
+        "OrbitalRotation",
+    )
+
+    for idx in range(tape.num_params):
+        op, _ = tape.get_operation(idx)
+        if op.name not in supported_ops:
+            raise ValueError(
+                f"The operation {op.name} is currently not supported for the parameter-shift "
+                f"Hessian. Only two-term parameter shift rules are currently supported."
+            )
+
+    _gradient_analysis(tape)
+    diff_methods = tape._grad_method_validation("analytic")  # pylint: disable=protected-access
+
+    return compute_hessian_tapes(tape, diff_methods, f0)
diff --git a/pennylane/math/single_dispatch.py b/pennylane/math/single_dispatch.py
index f2b86e49e..23abda165 100644
--- a/pennylane/math/single_dispatch.py
+++ b/pennylane/math/single_dispatch.py
@@ -287,6 +287,14 @@ def _scatter_element_add_tf(tensor, index, value):
 ar.register_function("tensorflow", "scatter_element_add", _scatter_element_add_tf)
 
 
+def _transpose_tf(a, axes=None):
+    import tensorflow as tf
+
+    return tf.transpose(a, perm=axes)
+
+
+ar.register_function("tensorflow", "transpose", _transpose_tf)
+
 # -------------------------------- Torch --------------------------------- #
 
 ar.autoray._FUNC_ALIASES["torch", "unstack"] = "unbind"


Test patch:
diff --git a/tests/gradients/test_param_shift_hessian.py b/tests/gradients/test_param_shift_hessian.py
new file mode 100644
index 000000000..14610311b
--- /dev/null
+++ b/tests/gradients/test_param_shift_hessian.py
@@ -0,0 +1,693 @@
+# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for the gradients.param_shift_hessian module."""
+
+import pytest
+
+import pennylane as qml
+from pennylane import numpy as np
+
+
+class TestParameterShiftHessian:
+    """Test the general functionality of the param_shift_hessian method
+    on the default interface (autograd)"""
+
+    def test_single_two_term_gate(self):
+        """Test that the correct hessian is calculated for a QNode with single RX operator
+        and single expectation value output (0d -> 0d)"""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="parameter-shift", max_diff=2)
+        def circuit(x):
+            qml.RX(x, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0))
+
+        x = np.array(0.1, requires_grad=True)
+
+        expected = qml.jacobian(qml.grad(circuit))(x)
+        hessian = qml.gradients.param_shift_hessian(circuit)(x)
+
+        assert np.allclose(expected, hessian)
+
+    def test_single_two_term_gate_vector_output(self):
+        """Test that the correct hessian is calculated for a QNode with single RY operator
+        and probabilies as output (0d -> 1d)"""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="parameter-shift", max_diff=2)
+        def circuit(x):
+            qml.RY(x, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.probs(wires=[0, 1])
+
+        x = np.array(0.1, requires_grad=True)
+
+        expected = qml.jacobian(qml.jacobian(circuit))(x)
+        hessian = qml.gradients.param_shift_hessian(circuit)(x)
+
+        assert np.allclose(expected, hessian)
+
+    def test_multiple_two_term_gates(self):
+        """Test that the correct hessian is calculated for a QNode with two rotation operators
+        and one expectation value output (1d -> 0d)"""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="parameter-shift", max_diff=2)
+        def circuit(x):
+            qml.RX(x[0], wires=0)
+            qml.RY(x[1], wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(1))
+
+        x = np.array([0.1, 0.2], requires_grad=True)
+
+        expected = qml.jacobian(qml.jacobian(circuit))(x)
+        hessian = qml.gradients.param_shift_hessian(circuit)(x)
+
+        assert np.allclose(expected, hessian)
+
+    def test_multiple_two_term_gates_vector_output(self):
+        """Test that the correct hessian is calculated for a QNode with two rotation operators
+        and probabilities output (1d -> 1d)"""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="parameter-shift", max_diff=2)
+        def circuit(x):
+            qml.RX(x[0], wires=0)
+            qml.RY(x[1], wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.probs(wires=1)
+
+        x = np.array([0.1, 0.2], requires_grad=True)
+
+        expected = qml.jacobian(qml.jacobian(circuit))(x)
+        hessian = qml.gradients.param_shift_hessian(circuit)(x)
+
+        assert np.allclose(expected, hessian)
+
+    def test_quantum_hessian_shape_vector_input_vector_output(self):
+        """Test that the purely "quantum" hessian has the correct shape (1d -> 1d)"""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="parameter-shift", max_diff=2)
+        def circuit(x):
+            qml.RX(x[0], wires=0)
+            qml.RY(x[1], wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RZ(x[2], wires=1)
+            qml.Rot(x[0], x[1], x[2], wires=1)
+            return qml.probs(wires=0)
+
+        x = np.array([0.1, 0.2, 0.3], requires_grad=True)
+        shape = (2, 6, 6)  # (num_output_vals, num_gate_args, num_gate_args)
+
+        hessian = qml.gradients.param_shift_hessian(circuit, hybrid=False)(x)
+
+        assert qml.math.shape(hessian) == shape
+
+    def test_multiple_two_term_gates_reusing_parameters(self):
+        """Test that the correct hessian is calculated when reusing parameters (1d -> 1d)"""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="parameter-shift", max_diff=2)
+        def circuit(x):
+            qml.RX(x[0], wires=0)
+            qml.RY(x[1], wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RZ(x[2], wires=1)
+            qml.Rot(x[0], x[1], x[2], wires=1)
+            return qml.probs(wires=0)
+
+        x = np.array([0.1, 0.2, 0.3], requires_grad=True)
+
+        expected = qml.jacobian(qml.jacobian(circuit))(x)
+        hessian = qml.gradients.param_shift_hessian(circuit)(x)
+
+        assert np.allclose(expected, hessian)
+
+    def test_multiple_two_term_gates_classical_processing(self):
+        """Test that the correct hessian is calculated when manipulating parameters (1d -> 1d)"""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="parameter-shift", max_diff=2)
+        def circuit(x):
+            qml.RX(x[0] + x[1] + x[2], wires=0)
+            qml.RY(x[1] - x[0] + 3 * x[2], wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RZ(x[2] / x[0] - x[1], wires=1)
+            return qml.probs(wires=0)
+
+        x = np.array([0.1, 0.2, 0.3], requires_grad=True)
+
+        expected = qml.jacobian(qml.jacobian(circuit))(x)
+        hessian = qml.gradients.param_shift_hessian(circuit)(x)
+
+        assert np.allclose(expected, hessian)
+
+    def test_multiple_two_term_gates_matrix_output(self):
+        """Test that the correct hessian is calculated for higher dimensional QNode outputs
+        (1d -> 2d)"""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="parameter-shift", max_diff=2)
+        def circuit(x):
+            qml.RX(x[0], wires=0)
+            qml.RY(x[1], wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.probs(wires=0), qml.probs(wires=1)
+
+        x = np.ones([2], requires_grad=True)
+
+        expected = qml.jacobian(qml.jacobian(circuit))(x)
+        hessian = qml.gradients.param_shift_hessian(circuit)(x)
+
+        assert np.allclose(expected, hessian)
+
+    def test_multiple_two_term_gates_matrix_input(self):
+        """Test that the correct hessian is calculated for higher dimensional cl. jacobians
+        (2d -> 2d)"""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="parameter-shift", max_diff=2)
+        def circuit(x):
+            qml.RX(x[0, 0], wires=0)
+            qml.RY(x[0, 1], wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RX(x[0, 2], wires=0)
+            qml.RY(x[0, 0], wires=0)
+            return qml.probs(wires=0), qml.probs(wires=1)
+
+        x = np.ones([1, 3], requires_grad=True)
+
+        expected = qml.jacobian(qml.jacobian(circuit))(x)
+        hessian = qml.gradients.param_shift_hessian(circuit)(x)
+
+        assert np.allclose(expected, hessian)
+
+    @pytest.mark.slow
+    def test_multiple_qnode_arguments_scalar(self):
+        """Test that the correct Hessian is calculated with multiple QNode arguments (0D->1D)"""
+        jax = pytest.importorskip("jax")
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="backprop", interface="jax")
+        def circuit(x, y, z):
+            qml.RX(x, wires=0)
+            qml.RY(y, wires=1)
+            qml.RZ(z, wires=1)
+            qml.RY(y, wires=0)
+            qml.RX(x, wires=0)
+            return qml.probs(wires=[0, 1])
+
+        x = np.array(0.1, requires_grad=True)
+        y = np.array(0.5, requires_grad=True)
+        z = np.array(0.3, requires_grad=True)
+
+        expected = tuple(
+            jax.jacobian(jax.jacobian(circuit, argnums=i), argnums=i)(
+                jax.numpy.array(x), jax.numpy.array(y), jax.numpy.array(z)
+            )
+            for i in range(3)
+        )
+        circuit.interface = "autograd"
+        hessian = qml.gradients.param_shift_hessian(circuit)(x, y, z)
+
+        assert np.allclose(expected, hessian)
+
+    @pytest.mark.slow
+    def test_multiple_qnode_arguments_vector(self):
+        """Test that the correct Hessian is calculated with multiple QNode arguments (1D->1D)"""
+        jax = pytest.importorskip("jax")
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="backprop", interface="jax")
+        def circuit(x, y, z):
+            qml.RX(x[0], wires=1)
+            qml.RY(y[0], wires=0)
+            qml.RZ(z[0] + z[1], wires=1)
+            qml.RY(y[1], wires=1)
+            qml.RX(x[1], wires=0)
+            return qml.probs(wires=[0, 1])
+
+        x = np.array([0.1, 0.3], requires_grad=True)
+        y = np.array([0.5, 0.7], requires_grad=True)
+        z = np.array([0.3, 0.2], requires_grad=True)
+
+        expected = tuple(
+            jax.jacobian(jax.jacobian(circuit, argnums=i), argnums=i)(
+                jax.numpy.array(x), jax.numpy.array(y), jax.numpy.array(z)
+            )
+            for i in range(3)
+        )
+        circuit.interface = "autograd"
+        hessian = qml.gradients.param_shift_hessian(circuit)(x, y, z)
+
+        assert np.allclose(expected, hessian)
+
+    @pytest.mark.slow
+    def test_multiple_qnode_arguments_matrix(self):
+        """Test that the correct Hessian is calculated with multiple QNode arguments (2D->1D)"""
+        jax = pytest.importorskip("jax")
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="backprop", interface="jax")
+        def circuit(x, y, z):
+            qml.RX(x[0, 0], wires=0)
+            qml.RY(y[0, 0], wires=1)
+            qml.RZ(z[0, 0] + z[1, 1], wires=1)
+            qml.RY(y[1, 0], wires=0)
+            qml.RX(x[1, 0], wires=1)
+            return qml.probs(wires=[0, 1])
+
+        x = np.array([[0.1, 0.3], [0.2, 0.4]], requires_grad=True)
+        y = np.array([[0.5, 0.7], [0.2, 0.4]], requires_grad=True)
+        z = np.array([[0.3, 0.2], [0.2, 0.4]], requires_grad=True)
+
+        expected = tuple(
+            jax.jacobian(jax.jacobian(circuit, argnums=i), argnums=i)(
+                jax.numpy.array(x), jax.numpy.array(y), jax.numpy.array(z)
+            )
+            for i in range(3)
+        )
+        circuit.interface = "autograd"
+        hessian = qml.gradients.param_shift_hessian(circuit)(x, y, z)
+
+        assert np.allclose(expected, hessian)
+
+    def test_multiple_qnode_arguments_mixed(self):
+        """Test that the correct Hessian is calculated with multiple mixed-shape QNode arguments"""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="parameter-shift", max_diff=2)
+        def circuit(x, y, z):
+            qml.RX(x, wires=0)
+            qml.RY(z[0] + z[1], wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RX(y[1, 0], wires=0)
+            qml.RY(y[0, 1], wires=0)
+            return qml.probs(wires=0), qml.probs(wires=1)
+
+        x = np.array(0.1, requires_grad=True)
+        y = np.array([[0.5, 0.6], [0.2, 0.1]], requires_grad=True)
+        z = np.array([0.3, 0.4], requires_grad=True)
+
+        expected = tuple(
+            qml.jacobian(qml.jacobian(circuit, argnum=i), argnum=i)(x, y, z) for i in range(3)
+        )
+        hessian = qml.gradients.param_shift_hessian(circuit)(x, y, z)
+
+        assert all(np.allclose(expected[i], hessian[i]) for i in range(3))
+
+    def test_hessian_transform_is_differentiable(self):
+        """Test that the 3rd derivate can be calculated via auto-differentiation (1d -> 1d)"""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="parameter-shift", max_diff=3)
+        def circuit(x):
+            qml.RX(x[1], wires=0)
+            qml.RY(x[0], wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.probs(wires=[0, 1])
+
+        x = np.array([0.1, 0.2], requires_grad=True)
+
+        expected = qml.jacobian(qml.jacobian(qml.jacobian(circuit)))(x)
+        derivative = qml.jacobian(qml.gradients.param_shift_hessian(circuit))(x)
+
+        assert np.allclose(expected, derivative)
+
+    # Some bounds on the efficiency (device executions) of the hessian for 2-term shift rules:
+    # - < jacobian(jacobian())
+    # - <= 2^d * (m+d-1)C(d)      see arXiv:2008.06517 p. 4
+    # - <= 3^m                    see arXiv:2008.06517 p. 4
+    # here d=2 is the derivative order, m is the number of variational parameters (w.r.t. gate args)
+
+    def test_fewer_device_invocations_scalar_input(self):
+        """Test that the hessian invokes less hardware executions than double differentiation
+        (0d -> 0d)"""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="parameter-shift", max_diff=2)
+        def circuit(x):
+            qml.RX(x, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(1))
+
+        x = np.array(0.1, requires_grad=True)
+
+        with qml.Tracker(dev) as tracker:
+            hessian = qml.gradients.param_shift_hessian(circuit)(x)
+            hessian_qruns = tracker.totals["executions"]
+            expected = qml.jacobian(qml.jacobian(circuit))(x)
+            jacobian_qruns = tracker.totals["executions"] - hessian_qruns
+
+        assert np.allclose(hessian, expected)
+        assert hessian_qruns < jacobian_qruns
+        assert hessian_qruns <= 2 ** 2 * 1  # 1 = (1+2-1)C(2)
+        assert hessian_qruns <= 3 ** 1
+
+    def test_fewer_device_invocations_vector_input(self):
+        """Test that the hessian invokes less hardware executions than double differentiation
+        (1d -> 0d)"""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="parameter-shift", max_diff=2)
+        def circuit(x):
+            qml.RX(x[0], wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RY(x[1], wires=0)
+            return qml.expval(qml.PauliZ(1))
+
+        x = np.array([0.1, 0.2], requires_grad=True)
+
+        with qml.Tracker(dev) as tracker:
+            hessian = qml.gradients.param_shift_hessian(circuit)(x)
+            hessian_qruns = tracker.totals["executions"]
+            expected = qml.jacobian(qml.jacobian(circuit))(x)
+            jacobian_qruns = tracker.totals["executions"] - hessian_qruns
+
+        assert np.allclose(hessian, expected)
+        assert hessian_qruns < jacobian_qruns
+        assert hessian_qruns <= 2 ** 2 * 3  # 3 = (2+2-1)C(2)
+        assert hessian_qruns <= 3 ** 2
+
+    def test_fewer_device_invocations_vector_output(self):
+        """Test that the hessian invokes less hardware executions than double differentiation
+        (1d -> 1d)"""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="parameter-shift", max_diff=2)
+        def circuit(x):
+            qml.RX(x[0], wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RY(x[1], wires=0)
+            qml.RZ(x[2], wires=1)
+            return qml.probs(wires=[0, 1])
+
+        x = np.array([0.1, 0.2, 0.3], requires_grad=True)
+
+        with qml.Tracker(dev) as tracker:
+            hessian = qml.gradients.param_shift_hessian(circuit)(x)
+            hessian_qruns = tracker.totals["executions"]
+            expected = qml.jacobian(qml.jacobian(circuit))(x)
+            jacobian_qruns = tracker.totals["executions"] - hessian_qruns
+
+        assert np.allclose(hessian, expected)
+        assert hessian_qruns < jacobian_qruns
+        assert hessian_qruns <= 2 ** 2 * 6  # 6 = (3+2-1)C(2)
+        assert hessian_qruns <= 3 ** 3
+
+    def test_error_unsupported_operation(self):
+        """Test that the correct error is thrown for unsopperted operations"""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="parameter-shift", max_diff=2)
+        def circuit(x):
+            qml.RX(x[0], wires=0)
+            qml.RY(x[1], wires=0)
+            qml.CRZ(x[2], wires=[0, 1])
+            return qml.probs(wires=1)
+
+        x = np.array([0.1, 0.2, 0.3], requires_grad=True)
+
+        with pytest.raises(ValueError, match=r"The operation .+ is currently not supported"):
+            qml.gradients.param_shift_hessian(circuit)(x)
+
+    def test_error_unsupported_variance_measurement(self):
+        """Test that the correct error is thrown for variance measurements"""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="parameter-shift", max_diff=2)
+        def circuit(x):
+            qml.RX(x[0], wires=0)
+            qml.RY(x[1], wires=0)
+            qml.CRZ(x[2], wires=[0, 1])
+            return qml.var(qml.PauliZ(1))
+
+        x = np.array([0.1, 0.2, 0.3], requires_grad=True)
+
+        with pytest.raises(
+            ValueError,
+            match="Computing the Hessian of circuits that return variances is currently not supported.",
+        ):
+            qml.gradients.param_shift_hessian(circuit)(x)
+
+    def test_error_unsupported_state_measurement(self):
+        """Test that the correct error is thrown for state measurements"""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="parameter-shift", max_diff=2)
+        def circuit(x):
+            qml.RX(x[0], wires=0)
+            qml.RY(x[1], wires=0)
+            qml.CRZ(x[2], wires=[0, 1])
+            return qml.state()
+
+        x = np.array([0.1, 0.2, 0.3], requires_grad=True)
+
+        with pytest.raises(
+            ValueError,
+            match="Computing the Hessian of circuits that return the state is not supported.",
+        ):
+            qml.gradients.param_shift_hessian(circuit)(x)
+
+    def test_no_error_nondifferentiable_unsupported_operation(self):
+        """Test that no error is thrown for operations that are not marked differentiable"""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="parameter-shift", max_diff=2)
+        def circuit(x, y, z):
+            qml.RX(x, wires=0)
+            qml.RY(y, wires=0)
+            qml.CRZ(z, wires=[0, 1])
+            return qml.probs(wires=1)
+
+        x = np.array(0.1, requires_grad=True)
+        y = np.array(0.2, requires_grad=True)
+        z = np.array(0.3, requires_grad=False)
+
+        qml.gradients.param_shift_hessian(circuit)(x, y, z)
+
+    def test_no_trainable_parameters(self):
+        """Test that the correct ouput is generated in the absence of any trainable parameters"""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="parameter-shift", max_diff=2)
+        def circuit(x):
+            qml.RX(x[0], wires=0)
+            qml.RY(x[1], wires=0)
+            qml.CRZ(x[2], wires=[0, 1])
+            return qml.probs(wires=1)
+
+        x = np.array([0.1, 0.2, 0.3], requires_grad=False)
+
+        expected = qml.jacobian(qml.jacobian(circuit))(x)
+        hessian = qml.gradients.param_shift_hessian(circuit)(x)
+
+        assert np.allclose(expected, hessian)
+
+    def test_f0_argument(self):
+        """Test that we can provide the results of a QNode to save on quantum invocations"""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="parameter-shift", max_diff=2)
+        def circuit(x):
+            qml.RX(x[0], wires=0)
+            qml.RY(x[1], wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.probs(wires=1)
+
+        x = np.array([0.1, 0.2], requires_grad=True)
+
+        res = circuit(x)
+
+        with qml.Tracker(dev) as tracker:
+            hessian1 = qml.gradients.param_shift_hessian(circuit, f0=res)(x)
+            qruns1 = tracker.totals["executions"]
+            hessian2 = qml.gradients.param_shift_hessian(circuit)(x)
+            qruns2 = tracker.totals["executions"] - qruns1
+
+        assert np.allclose(hessian1, hessian2)
+        assert qruns1 < qruns2
+
+
+class TestInterfaces:
+    """Test the param_shift_hessian method on different interfaces"""
+
+    def test_hessian_transform_with_torch(self):
+        """Test that the Hessian transform can be used with Torch (1d -> 1d)"""
+        torch = pytest.importorskip("torch")
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RX(x[1], wires=0)
+            qml.RY(x[0], wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.probs(wires=[0, 1])
+
+        x_np = np.array([0.1, 0.2], requires_grad=True)
+        x_torch = torch.tensor([0.1, 0.2], dtype=torch.float64, requires_grad=True)
+
+        expected = qml.jacobian(qml.jacobian(circuit))(x_np)
+        circuit.interface = "torch"
+        hess = qml.gradients.param_shift_hessian(circuit)(x_torch)[0]
+
+        assert np.allclose(expected, hess.detach())
+
+    def test_hessian_transform_is_differentiable_torch(self):
+        """Test that the 3rd derivate can be calculated via auto-differentiation in Torch
+        (1d -> 1d)"""
+        torch = pytest.importorskip("torch")
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="parameter-shift", max_diff=3)
+        def circuit(x):
+            qml.RX(x[1], wires=0)
+            qml.RY(x[0], wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.probs(wires=[0, 1])
+
+        x = np.array([0.1, 0.2], requires_grad=True)
+        x_torch = torch.tensor([0.1, 0.2], dtype=torch.float64, requires_grad=True)
+
+        expected = qml.jacobian(qml.jacobian(qml.jacobian(circuit)))(x)
+        circuit.interface = "torch"
+        jacobian_fn = torch.autograd.functional.jacobian
+        torch_deriv = jacobian_fn(qml.gradients.param_shift_hessian(circuit), x_torch)[0]
+
+        assert np.allclose(expected, torch_deriv)
+
+    @pytest.mark.slow
+    def test_hessian_transform_with_jax(self):
+        """Test that the Hessian transform can be used with JAX (1d -> 1d)"""
+        jax = pytest.importorskip("jax")
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RX(x[1], wires=0)
+            qml.RY(x[0], wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.probs(wires=[0, 1])
+
+        x_np = np.array([0.1, 0.2], requires_grad=True)
+        x_jax = jax.numpy.array([0.1, 0.2])
+
+        expected = qml.jacobian(qml.jacobian(circuit))(x_np)
+        circuit.interface = "jax"
+        hess = qml.gradients.param_shift_hessian(circuit)(x_jax)
+
+        assert np.allclose(expected, hess)
+
+    @pytest.mark.slow
+    def test_hessian_transform_is_differentiable_jax(self):
+        """Test that the 3rd derivate can be calculated via auto-differentiation in JAX
+        (1d -> 1d)"""
+        jax = pytest.importorskip("jax")
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="backprop", max_diff=3)
+        def circuit(x):
+            qml.RX(x[1], wires=0)
+            qml.RY(x[0], wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.probs(wires=[0, 1])
+
+        x = np.array([0.1, 0.2], requires_grad=True)
+        x_jax = jax.numpy.array([0.1, 0.2])
+
+        expected = qml.jacobian(qml.jacobian(qml.jacobian(circuit)))(x)
+        circuit.interface = "jax"
+        jax_deriv = jax.jacobian(qml.gradients.param_shift_hessian(circuit))(x_jax)
+
+        assert np.allclose(expected, jax_deriv)
+
+    @pytest.mark.slow
+    def test_hessian_transform_with_tensorflow(self):
+        """Test that the Hessian transform can be used with TensorFlow (1d -> 1d)"""
+        tf = pytest.importorskip("tensorflow")
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RX(x[1], wires=0)
+            qml.RY(x[0], wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.probs(wires=[0, 1])
+
+        x_np = np.array([0.1, 0.2], requires_grad=True)
+        x_tf = tf.Variable([0.1, 0.2], dtype=tf.float64)
+
+        expected = qml.jacobian(qml.jacobian(circuit))(x_np)
+        circuit.interface = "tf"
+        with tf.GradientTape():
+            hess = qml.gradients.param_shift_hessian(circuit)(x_tf)[0]
+
+        assert np.allclose(expected, hess)
+
+    @pytest.mark.slow
+    def test_hessian_transform_is_differentiable_tensorflow(self):
+        """Test that the 3rd derivate can be calculated via auto-differentiation in Tensorflow
+        (1d -> 1d)"""
+        tf = pytest.importorskip("tensorflow")
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="parameter-shift", max_diff=3)
+        def circuit(x):
+            qml.RX(x[1], wires=0)
+            qml.RY(x[0], wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.probs(wires=[0, 1])
+
+        x = np.array([0.1, 0.2], requires_grad=True)
+        x_tf = tf.Variable([0.1, 0.2], dtype=tf.float64)
+
+        expected = qml.jacobian(qml.jacobian(qml.jacobian(circuit)))(x)
+        circuit.interface = "tf"
+        with tf.GradientTape() as tf_tape:
+            hessian = qml.gradients.param_shift_hessian(circuit)(x_tf)[0]
+        tensorflow_deriv = tf_tape.jacobian(hessian, x_tf)
+
+        assert np.allclose(expected, tensorflow_deriv)