Repository: PennyLaneAI/pennylane
Pull request #1476 (2021-07-26 14:16:00)

Differentiable finite-difference gradient transform
**Context:** As part of the roadmap for supporting differentiable batch execution, gradient logic will be moved out of the subclasses and into a module of pure functions. This is the first such PR; here, we move the finite-difference logic out of `JacobianTape` and into a new `gradients` package.

**Description of the Change:**

- Adds a function `qml.gradients.finite_diff_stencil`, for generating the coefficients and shifts required for finite-difference rules of various order, accuracy, and form.

- Adds a function `qml.gradients.finite_diff` for generating the finite-difference tapes of an input tape. The transform is fully differentiable, and may be differentiated to get higher-order derivatives.

**Benefits:** The finite-difference logic is now much more user and dev accessible, and supports higher order derivatives.

**Possible Drawbacks:**

Currently, when executing a device, the output is always cast to a tensor. For example, `expval(), expval() -> tensor[2]`, and `probs[0], probs[1] -> tensor[2, 2]`. However, this results in _ragged_ arrays in cases such as `expval(), probs()`.

This has two drawbacks:

- NumPy is deprecating ragged arrays shortly (!!)
- Some ML frameworks don't support ragged arrays

In the current codebase, we get around the latter issue by always `hstack`-ing the ragged arrays when we compute the Jacobian. E.g., for N input parameters and output `expval(), probs([0, 1]) -> tensor[[1], [4]]`, the Jacobian will be returned as a `[5, N]` array (the 1 and 4 dimensions of the output have been concatenated).

However, long term, we should move to treating the output as a tuple  `expval(), probs([0, 1]) -> tuple[tensor[1], tensor[4]]`. This avoids the issues listed in the bullet points above, and also makes the jacobian computation easier; we simply return `tuple[tensor[1, N], tensor[4, N]]`.

**Related GitHub Issues:** n/a


josh146: [ch7194]
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1476?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1476](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1476?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (c3a5b41) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/9a032501cb4162410cafec8c4d907a6b4b8d2414?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (9a03250) will **increase** coverage by `0.01%`.
> The diff coverage is `100.00%`.

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

```diff
@@            Coverage Diff             @@
##           master    #1476      +/-   ##
==========================================
+ Coverage   98.30%   98.31%   +0.01%     
==========================================
  Files         170      172       +2     
  Lines       12483    12574      +91     
==========================================
+ Hits        12271    12362      +91     
  Misses        212      212              
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1476?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1476/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19faW5pdF9fLnB5) | `98.59% <100.00%> (+0.02%)` | :arrow_up: |
| [pennylane/gradients/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1476/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9fX2luaXRfXy5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/gradients/finite\_difference.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1476/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9maW5pdGVfZGlmZmVyZW5jZS5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/math/single\_dispatch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1476/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvc2luZ2xlX2Rpc3BhdGNoLnB5) | `99.43% <100.00%> (+0.01%)` | :arrow_up: |

------

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

josh146: > How user-facing is this transform? Is it something that would also be nice to have a QNode transform wrapper for so users don't have to work directly with tapes?

At _some point_ I would like to make it more user-friendly, but at the moment it is purely a port of `JacobianTape.jacobian` from a method to a transform. The plan is:

1. The new, differentiable `batch_execute` will use this transform internally when computing gradients during backpropagation.
2. Later, wrap this up as a qnode transform so it can also be used by users. However, this is blocked by us putting the QNode transform API on hold!

> While there is a parameter n to set the order of the gradient, I mostly noticed n=1 and a few n=2 cases in the tests. It may be valuable to show a larger example; also, the transform can be applied to something multiple times, right? If so would be cool to see an example of that as well.

I have to admit, I planned to only port the code in `JacobianTape.jacobian` as-is to reduce my workload. It currently _only_ supports first derivatives, and only 2nd-order center and 1st-order forward.

However.... I went a bit down a rabbit hole while porting it, and instead implemented support for arbitrary finite-diff rules 🙁 

josh146: > but I requested changes because I was often quite confused, and ideally we want to make the code base self-explanatory.

@mariaschuld these were good catches. I've updated the PR accordingly, hopefully it should be easier to follow now!

Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index ff849798c..c5bc8f887 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -2,6 +2,31 @@
 
 <h3>New features since last release</h3>
 
+* A new gradients module `qml.gradients` has been added, which provides
+  differentiable quantum gradient transforms.
+  [(#1476)](https://github.com/PennyLaneAI/pennylane/pull/1476)
+
+  Available quantum gradient transforms include:
+
+  - `qml.gradients.finite_diff`
+
+  For example,
+
+  ```pycon
+  >>> with qml.tape.QuantumTape() as tape:
+  ...     qml.RX(params[0], wires=0)
+  ...     qml.RY(params[1], wires=0)
+  ...     qml.RX(params[2], wires=0)
+  ...     qml.expval(qml.PauliZ(0))
+  ...     qml.var(qml.PauliZ(0))
+  >>> tape.trainable_params = {0, 1, 2}
+  >>> gradient_tapes, fn = qml.gradients.finite_diff(tape)
+  >>> res = dev.batch_execute(gradient_tapes)
+  >>> fn(res)
+  [[-0.38751721 -0.18884787 -0.38355704]
+   [ 0.69916862  0.34072424  0.69202359]]
+  ```
+
 * A new quantum function transform has been added to push commuting
   single-qubit gates through controlled operations.
   [(#1464)](https://github.com/PennyLaneAI/pennylane/pull/1464)
diff --git a/pennylane/__init__.py b/pennylane/__init__.py
index bcb19c199..3da94819a 100644
--- a/pennylane/__init__.py
+++ b/pennylane/__init__.py
@@ -25,6 +25,7 @@ from pennylane.queuing import apply, QueuingContext
 
 import pennylane.init
 import pennylane.fourier
+import pennylane.gradients
 import pennylane.kernels
 import pennylane.math
 import pennylane.operation
diff --git a/pennylane/gradients/__init__.py b/pennylane/gradients/__init__.py
new file mode 100644
index 000000000..3cbfaf034
--- /dev/null
+++ b/pennylane/gradients/__init__.py
@@ -0,0 +1,19 @@
+# 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 subpackage contains quantum gradient transforms."""
+import pennylane as qml
+
+from . import finite_difference
+
+from .finite_difference import finite_diff, finite_diff_coeffs, generate_shifted_tapes
diff --git a/pennylane/gradients/finite_difference.py b/pennylane/gradients/finite_difference.py
new file mode 100644
index 000000000..557ed7035
--- /dev/null
+++ b/pennylane/gradients/finite_difference.py
@@ -0,0 +1,306 @@
+# 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 finite-difference gradient
+of a quantum tape.
+"""
+# pylint: disable=protected-access,too-many-arguments
+import numpy as np
+from scipy.special import factorial
+
+import pennylane as qml
+
+
+def finite_diff_coeffs(n, approx_order, strategy):
+    r"""Generate the finite difference shift values and corresponding
+    term coefficients for a given derivative order, approximation accuracy,
+    and strategy.
+
+    Args:
+        n (int): Positive integer specifying the order of the derivative. For example, ``n=1``
+            corresponds to the first derivative, ``n=2`` the second derivative, etc.
+        approx_order (int): Positive integer referring to the approximation order of the
+            returned coefficients, e.g., ``approx_order=1`` corresponds to the
+            first-order approximation to the derivative.
+        strategy (str): One of ``"forward"``, ``"center"``, or ``"backward"``.
+            For the ``"forward"`` strategy, the finite-difference shifts occur at the points
+            :math:`x_0, x_0+h, x_0+2h,\dots`, where :math:`h` is some small
+            step size. The ``"backwards"`` strategy is similar, but in
+            reverse: :math:`x_0, x_0-h, x_0-2h, \dots`. Finally, the
+            ``"center"`` strategy results in shifts symmetric around the
+            unshifted point: :math:`\dots, x_0-2h, x_0-h, x_0, x_0+h, x_0+2h,\dots`.
+
+    Returns:
+        array[float]: A ``(2, N)`` array. The first row corresponds to the
+        coefficients, and the second row corresponds to the shifts.
+
+    **Example**
+
+    >>> finite_diff_coeffs(n=1, approx_order=1, strategy="forward")
+    array([[-1.,  1.],
+           [ 0.,  1.]])
+
+    For example, this results in the linear combination:
+
+    .. math:: \frac{-y(x_0) + y(x_0 + h)}{h}
+
+    where :math:`h` is the finite-difference step size.
+
+    More examples:
+
+    >>> finite_diff_coeffs(n=1, approx_order=2, strategy="center")
+    array([[-0.5,  0.5],
+           [-1. ,  1. ]])
+    >>> finite_diff_coeffs(n=2, approx_order=2, strategy="center")
+    array([[-2.,  1.,  1.],
+           [ 0., -1.,  1.]])
+
+    **Details**
+
+    Consider a function :math:`y(x)`. We wish to approximate the :math:`n`-th
+    derivative at point :math:`x_0`, :math:`y^{(n)}(x_0)`, by sampling the function
+    at :math:`N<n` distinct points:
+
+    .. math:: y^{(n)}(x_0) \approx \sum_{i=1}^N c_i y(x_i)
+
+    where :math:`c_i` are coefficients, and :math:`x_i=x_0 + s_i` are the points we sample
+    the function at.
+
+    Consider the Taylor expansion of :math:`y(x_i)` around the point :math:`x_0`:
+
+    .. math::
+
+        y^{(n)}(x_0) \approx \sum_{i=1}^N c_i y(x_i)
+            &= \sum_{i=1}^N c_i \left[ y(x_0) + y'(x_0)(x_i-x_0) + \frac{1}{2} y''(x_0)(x_i-x_0)^2 + \cdots \right]\\
+            & = \sum_{j=0}^m y^{(j)}(x_0) \left[\sum_{i=1}^N \frac{c_i s_i^j}{j!} + \mathcal{O}(s_i^m) \right],
+
+    where :math:`s_i = x_i-x_0`. For this approximation to be satisfied, we must therefore have
+
+    .. math::
+
+        \sum_{i=1}^N s_i^j c_i = \begin{cases} j!, &j=n\\ 0, & j\neq n\end{cases}.
+
+    Thus, to determine the coefficients :math:`c_i \in \{c_1, \dots, c_N\}` for particular
+    shift values :math:`s_i \in \{s_1, \dots, s_N\}` and derivative order :math:`n`,
+    we must solve this linear system of equations.
+    """
+    if n < 1 or not isinstance(n, int):
+        raise ValueError("Derivative order n must be a positive integer.")
+
+    if approx_order < 1 or not isinstance(approx_order, int):
+        raise ValueError("Approximation order must be a positive integer.")
+
+    num_points = approx_order + 2 * np.floor((n + 1) / 2) - 1
+    N = num_points + 1 if n % 2 == 0 else num_points
+
+    if strategy == "forward":
+        shifts = np.arange(N, dtype=np.float64)
+
+    elif strategy == "backward":
+        shifts = np.arange(-N + 1, 1, dtype=np.float64)
+
+    elif strategy == "center":
+        if approx_order % 2 != 0:
+            raise ValueError("Centered finite-difference requires an even order approximation.")
+
+        N = num_points // 2
+        shifts = np.arange(-N, N + 1, dtype=np.float64)
+
+    else:
+        raise ValueError(
+            f"Unknown strategy {strategy}. Must be one of 'forward', 'backward', 'center'."
+        )
+
+    # solve for the coefficients
+    A = shifts ** np.arange(len(shifts)).reshape(-1, 1)
+    b = np.zeros_like(shifts)
+    b[n] = factorial(n)
+    coeffs = np.linalg.solve(A, b)
+
+    coeffs_and_shifts = np.stack([coeffs, shifts])
+
+    # remove all small coefficients and shifts
+    coeffs_and_shifts[np.abs(coeffs_and_shifts) < 1e-10] = 0
+
+    # remove columns where the coefficients are 0
+    coeffs_and_shifts = coeffs_and_shifts[:, ~np.all(coeffs_and_shifts == 0, axis=0)]
+
+    # sort columns in ascending order according to abs(shift)
+    coeffs_and_shifts = coeffs_and_shifts[:, np.argsort(np.abs(coeffs_and_shifts)[1])]
+    return coeffs_and_shifts
+
+
+def generate_shifted_tapes(tape, idx, shifts, multipliers=None):
+    r"""Generate a list of tapes where the corresponding trainable parameter
+    index has been shifted by the values given.
+
+    Args:
+        tape (.QuantumTape): input quantum tape
+        idx (int): trainable parameter index to shift the parameter of
+        shifts (Sequence[float or int]): sequence of shift values
+        multipliers (Sequence[float or int]): Sequence of multiplier values to
+            scale the parameter by. If not provided, the parameter will
+            not be scaled.
+
+    Returns:
+        list[QuantumTape]: List of quantum tapes. Each tape has parameter
+        ``idx`` shifted by consecutive values of ``shift``. The length
+        of the returned list of tapes will match the length of ``shifts``.
+    """
+    params = tape.get_parameters()
+    tapes = []
+
+    for i, s in enumerate(shifts):
+        new_params = params.copy()
+        shifted_tape = tape.copy(copy_operations=True)
+
+        if multipliers is not None:
+            m = multipliers[i]
+            new_params[idx] = new_params[idx] * qml.math.convert_like(m, new_params[idx])
+
+        new_params[idx] = new_params[idx] + qml.math.convert_like(s, new_params[idx])
+        shifted_tape.set_parameters(new_params)
+        tapes.append(shifted_tape)
+
+    return tapes
+
+
+def finite_diff(tape, argnum=None, h=1e-7, approx_order=1, n=1, strategy="forward", f0=None):
+    r"""Generate the finite-difference tapes and postprocessing methods required
+    to compute the gradient of a gate parameter with respect to its outputs.
+
+    Args:
+        tape (.QuantumTape): quantum tape to differentiate
+        argnum (int or list[int] or None): Trainable parameter indices to differentiate
+            with respect to. If not provided, the derivatives with respect to all
+            trainable indices are returned.
+        h (float): finite difference method step size
+        approx_order (int): The approximation order of the finite-difference method to use.
+        n (int): compute the :math:`n`-th derivative
+        strategy (str): The strategy of the finite difference method. Must be one of
+            ``"forward"``, ``"center"``, or ``"backward"``.
+            For the ``"forward"`` strategy, the finite-difference shifts occur at the points
+            :math:`x_0, x_0+h, x_0+2h,\dots`, where :math:`h` is some small
+            stepsize. The ``"backwards"`` strategy is similar, but in
+            reverse: :math:`x_0, x_0-h, x_0-2h, \dots`. Finally, the
+            ``"center"`` strategy results in shifts symmetric around the
+            unshifted point: :math:`\dots, x_0-2h, x_0-h, x_0, x_0+h, x_0+2h,\dots`.
+        f0 (tensor_like[float] or None): Output of the evaluated input tape. If provided,
+            and the gradient recipe contains an unshifted term, this value is used,
+            saving a quantum evaluation.
+
+    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 evaluated tapes.
+
+    **Example**
+
+    >>> with qml.tape.QuantumTape() as tape:
+    ...     qml.RX(params[0], wires=0)
+    ...     qml.RY(params[1], wires=0)
+    ...     qml.RX(params[2], wires=0)
+    ...     qml.expval(qml.PauliZ(0))
+    ...     qml.var(qml.PauliZ(0))
+    >>> tape.trainable_params = {0, 1, 2}
+    >>> gradient_tapes, fn = qml.gradients.finite_diff.grad(tape)
+    >>> res = dev.batch_execute(gradient_tapes)
+    >>> fn(res)
+    [[-0.38751721 -0.18884787 -0.38355704]
+     [ 0.69916862  0.34072424  0.69202359]]
+
+    The output Jacobian matrix is of size ``(number_outputs, number_parameters)``.
+    """
+    # TODO: replace the JacobianTape._grad_method_validation
+    # functionality before deprecation.
+    diff_methods = tape._grad_method_validation("numeric")
+
+    if not tape.trainable_params or all(g == "0" for g in diff_methods):
+        # Either all parameters have grad method 0, or there are no trainable
+        # parameters.
+        return [], lambda x: np.zeros([tape.output_dim, len(tape.trainable_params)])
+
+    gradient_tapes = []
+    shapes = []
+    c0 = None
+
+    coeffs, shifts = finite_diff_coeffs(n=n, approx_order=approx_order, strategy=strategy)
+
+    if 0 in shifts:
+        # Finite difference formula includes a term with zero shift.
+
+        if f0 is None:
+            # Ensure that the unshifted tape is appended
+            # to the gradient tapes, if not already.
+            gradient_tapes.append(tape)
+
+        # Store the unshifted coefficient. We know that
+        # it will always be the first coefficient due to processing.
+        c0 = coeffs[0]
+        shifts = shifts[1:]
+        coeffs = coeffs[1:]
+
+    # TODO: replace the JacobianTape._choose_params_with_methods
+    # functionality before deprecation.
+    method_map = dict(tape._choose_params_with_methods(diff_methods, argnum))
+
+    for i, _ in enumerate(tape.trainable_params):
+        if i not in method_map or method_map[i] == "0":
+            # parameter has zero gradient
+            shapes.append(0)
+            continue
+
+        g_tapes = generate_shifted_tapes(tape, i, shifts * h)
+        gradient_tapes.extend(g_tapes)
+        shapes.append(len(g_tapes))
+
+    def processing_fn(results):
+        grads = []
+        start = 1 if c0 is not None and f0 is None else 0
+        r0 = f0 or results[0]
+
+        for s in shapes:
+
+            if s == 0:
+                # parameter has zero gradient
+                g = qml.math.zeros_like(results[0])
+                grads.append(g)
+                continue
+
+            res = results[start : start + s]
+            start = start + s
+
+            # compute the linear combination of results and coefficients
+            res = qml.math.stack(res)
+            g = sum([c * r for c, r in zip(coeffs, res)])
+
+            if c0 is not None:
+                # add on the unshifted term
+                g = g + c0 * r0
+
+            grads.append(g / (h ** n))
+
+        # The following is for backwards compatibility; currently,
+        # the device stacks multiple measurement arrays, even if not the same
+        # size, resulting in a ragged array.
+        # In the future, we might want to change this so that only tuples
+        # of arrays are returned.
+        for i, g in enumerate(grads):
+            g = qml.math.convert_like(g, res[0])
+            if hasattr(g, "dtype") and g.dtype is np.dtype("object"):
+                grads[i] = qml.math.hstack(g)
+
+        return qml.math.T(qml.math.stack(grads))
+
+    return gradient_tapes, processing_fn
diff --git a/pennylane/math/single_dispatch.py b/pennylane/math/single_dispatch.py
index 312daa9c2..2aea05ea9 100644
--- a/pennylane/math/single_dispatch.py
+++ b/pennylane/math/single_dispatch.py
@@ -40,6 +40,7 @@ ar.register_function("numpy", "coerce", lambda x: x)
 ar.register_function("numpy", "block_diag", lambda x: _scipy_block_diag(*x))
 ar.register_function("builtins", "block_diag", lambda x: _scipy_block_diag(*x))
 ar.register_function("numpy", "gather", lambda x, indices: x[np.array(indices)])
+ar.register_function("numpy", "unstack", list)
 
 
 def _scatter_element_add_numpy(tensor, index, value):
@@ -69,6 +70,7 @@ ar.register_function("autograd", "flatten", lambda x: x.flatten())
 ar.register_function("autograd", "coerce", lambda x: x)
 ar.register_function("autograd", "block_diag", lambda x: _scipy_block_diag(*x))
 ar.register_function("autograd", "gather", lambda x, indices: x[np.array(indices)])
+ar.register_function("autograd", "unstack", list)
 
 
 def _to_numpy_autograd(x):
@@ -230,6 +232,7 @@ ar.register_function("tensorflow", "scatter_element_add", _scatter_element_add_t
 
 # -------------------------------- Torch --------------------------------- #
 
+ar.autoray._FUNC_ALIASES["torch", "unstack"] = "unbind"
 
 ar.register_function("torch", "asarray", lambda x: _i("torch").as_tensor(x))
 ar.register_function("torch", "diag", lambda x, k=0: _i("torch").diag(x, diagonal=k))
@@ -357,3 +360,4 @@ ar.register_function(
     "scatter_element_add",
     lambda x, index, value: _i("jax").ops.index_add(x, tuple(index), value),
 )
+ar.register_function("jax", "unstack", list)


Test patch:
diff --git a/tests/gradients/test_finite_difference.py b/tests/gradients/test_finite_difference.py
new file mode 100644
index 000000000..0c4ce036b
--- /dev/null
+++ b/tests/gradients/test_finite_difference.py
@@ -0,0 +1,631 @@
+# 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.finite_difference module.
+"""
+import pytest
+
+from pennylane import numpy as np
+
+import pennylane as qml
+from pennylane.gradients import finite_diff, finite_diff_coeffs, generate_shifted_tapes
+
+
+class TestCoeffs:
+    """Tests for the finite_diff_coeffs function"""
+
+    def test_invalid_derivative_error(self):
+        """Test that an error is raised if n<1 or not an integer"""
+        with pytest.raises(ValueError, match="n must be a positive integer"):
+            finite_diff_coeffs(0, 1, 1)
+
+        with pytest.raises(ValueError, match="n must be a positive integer"):
+            finite_diff_coeffs(1.3, 1, 1)
+
+    def test_invalid_approx_order_error(self):
+        """Test that an error is raised if order < 1 or not an integer"""
+        with pytest.raises(ValueError, match="order must be a positive integer"):
+            finite_diff_coeffs(1, 0, 1)
+
+        with pytest.raises(ValueError, match="order must be a positive integer"):
+            finite_diff_coeffs(1, 1.7, 1)
+
+        with pytest.raises(ValueError, match="Centered finite-difference requires an even order"):
+            finite_diff_coeffs(1, 1, "center")
+
+    def test_invalid_strategy(self):
+        """Test that an error is raised if the strategy is not recognized"""
+        with pytest.raises(ValueError, match="Unknown strategy"):
+            finite_diff_coeffs(1, 1, 1)
+
+    def test_correct_forward_order1(self):
+        """Test that the correct forward order 1 method is returned"""
+        coeffs, shifts = finite_diff_coeffs(1, 1, "forward")
+        assert np.allclose(coeffs, [-1, 1])
+        assert np.allclose(shifts, [0, 1])
+
+    def test_correct_forward_order2(self):
+        """Test that the correct forward order 2 method is returned"""
+        coeffs, shifts = finite_diff_coeffs(1, 2, "forward")
+        assert np.allclose(coeffs, [-1.5, 2, -0.5])
+        assert np.allclose(shifts, [0, 1, 2])
+
+    def test_correct_center_order2(self):
+        """Test that the correct centered order 2 method is returned"""
+        coeffs, shifts = finite_diff_coeffs(1, 2, "center")
+        assert np.allclose(coeffs, [-0.5, 0.5])
+        assert np.allclose(shifts, [-1, 1])
+
+    def test_correct_backward_order1(self):
+        """Test that the correct backward order 1 method is returned"""
+        coeffs, shifts = finite_diff_coeffs(1, 1, "backward")
+        assert np.allclose(coeffs, [1, -1])
+        assert np.allclose(shifts, [0, -1])
+
+    def test_correct_second_derivative_forward_order1(self):
+        """Test that the correct forward order 1 method is returned"""
+        coeffs, shifts = finite_diff_coeffs(2, 1, "forward")
+        assert np.allclose(coeffs, [1, -2, 1])
+        assert np.allclose(shifts, [0, 1, 2])
+
+    def test_correct_second_derivative_center_order4(self):
+        """Test that the correct forward order 4 method is returned"""
+        coeffs, shifts = finite_diff_coeffs(2, 4, "center")
+        assert np.allclose(coeffs, [-2.5, 4 / 3, 4 / 3, -1 / 12, -1 / 12])
+        assert np.allclose(shifts, [0, -1, 1, -2, 2])
+
+
+class TestShiftedTapes:
+    """Tests for the generate_shifted_tapes function"""
+
+    def test_behaviour(self):
+        """Test that the function behaves as expected"""
+
+        with qml.tape.QuantumTape() as tape:
+            qml.PauliZ(0)
+            qml.RX(1.0, wires=0)
+            qml.CNOT(wires=[0, 2])
+            qml.Rot(2.0, 3.0, 4.0, wires=0)
+            qml.expval(qml.PauliZ(0))
+
+        tape.trainable_params = {0, 2}
+        shifts = [0.1, -0.2, 1.6]
+        res = generate_shifted_tapes(tape, 1, shifts=shifts)
+
+        assert len(res) == len(shifts)
+        assert res[0].get_parameters(trainable_only=False) == [1.0, 2.0, 3.1, 4.0]
+        assert res[1].get_parameters(trainable_only=False) == [1.0, 2.0, 2.8, 4.0]
+        assert res[2].get_parameters(trainable_only=False) == [1.0, 2.0, 4.6, 4.0]
+
+    def test_multipliers(self):
+        """Test that the function behaves as expected when multipliers are used"""
+
+        with qml.tape.JacobianTape() as tape:
+            qml.PauliZ(0)
+            qml.RX(1.0, wires=0)
+            qml.CNOT(wires=[0, 2])
+            qml.Rot(2.0, 3.0, 4.0, wires=0)
+            qml.expval(qml.PauliZ(0))
+
+        tape.trainable_params = {0, 2}
+        shifts = [0.3, 0.6]
+        multipliers = [0.2, 0.5]
+        res = generate_shifted_tapes(tape, 0, shifts=shifts, multipliers=multipliers)
+
+        assert len(res) == 2
+        assert res[0].get_parameters(trainable_only=False) == [0.2 * 1.0 + 0.3, 2.0, 3.0, 4.0]
+        assert res[1].get_parameters(trainable_only=False) == [0.5 * 1.0 + 0.6, 2.0, 3.0, 4.0]
+
+
+class TestFiniteDiff:
+    """Tests for the finite difference gradient transform"""
+
+    def test_non_differentiable_error(self):
+        """Test error raised if attempting to differentiate with
+        respect to a non-differentiable argument"""
+        psi = np.array([1, 0, 1, 0]) / np.sqrt(2)
+
+        with qml.tape.JacobianTape() as tape:
+            qml.QubitStateVector(psi, wires=[0, 1])
+            qml.RX(0.543, wires=[0])
+            qml.RY(-0.654, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            qml.probs(wires=[0, 1])
+
+        # by default all parameters are assumed to be trainable
+        with pytest.raises(
+            ValueError, match=r"Cannot differentiate with respect to parameter\(s\) {0}"
+        ):
+            finite_diff(tape)
+
+        # setting trainable parameters avoids this
+        tape.trainable_params = {1, 2}
+        dev = qml.device("default.qubit", wires=2)
+        tapes, fn = finite_diff(tape)
+
+        # For now, we must squeeze the results of the device execution, since
+        # qml.probs results in a nested result. Later, we will revisit device
+        # execution to avoid this issue.
+        res = fn(qml.math.squeeze(dev.batch_execute(tapes)))
+        assert res.shape == (4, 2)
+
+    def test_independent_parameter(self, mocker):
+        """Test that an independent parameter is skipped
+        during the Jacobian computation."""
+        spy = mocker.spy(qml.gradients.finite_difference, "generate_shifted_tapes")
+
+        with qml.tape.JacobianTape() as tape:
+            qml.RX(0.543, wires=[0])
+            qml.RY(-0.654, wires=[1])
+            qml.expval(qml.PauliZ(0))
+
+        dev = qml.device("default.qubit", wires=2)
+        tapes, fn = finite_diff(tape)
+        res = fn(dev.batch_execute(tapes))
+        assert res.shape == (1, 2)
+
+        assert len(spy.call_args_list) == 1
+
+        # only called for parameter 0
+        assert spy.call_args[0][0:2] == (tape, 0)
+
+    def test_no_trainable_parameters(self, mocker):
+        """Test that if the tape has no trainable parameters, no
+        subroutines are called and the returned Jacobian is empty"""
+        spy = mocker.spy(qml.gradients.finite_difference, "generate_shifted_tapes")
+
+        with qml.tape.JacobianTape() as tape:
+            qml.RX(0.543, wires=[0])
+            qml.RY(-0.654, wires=[1])
+            qml.expval(qml.PauliZ(0))
+
+        dev = qml.device("default.qubit", wires=2)
+        tape.trainable_params = {}
+
+        tapes, fn = finite_diff(tape)
+        res = fn(dev.batch_execute(tapes))
+        assert res.size == 0
+        assert np.all(res == np.array([[]]))
+
+        spy.assert_not_called()
+        assert len(tapes) == 0
+
+    def test_y0(self, mocker):
+        """Test that if first order finite differences is used, then
+        the tape is executed only once using the current parameter
+        values."""
+        dev = qml.device("default.qubit", wires=2)
+
+        with qml.tape.JacobianTape() as tape:
+            qml.RX(0.543, wires=[0])
+            qml.RY(-0.654, wires=[0])
+            qml.expval(qml.PauliZ(0))
+
+        tapes, fn = finite_diff(tape, approx_order=1)
+
+        # one tape per parameter, plus one global call
+        assert len(tapes) == tape.num_params + 1
+
+    def test_y0_provided(self):
+        """Test that if first order finite differences is used,
+        and the original tape output is provided, then
+        the tape is executed only once using the current parameter
+        values."""
+        dev = qml.device("default.qubit", wires=2)
+
+        with qml.tape.JacobianTape() as tape:
+            qml.RX(0.543, wires=[0])
+            qml.RY(-0.654, wires=[0])
+            qml.expval(qml.PauliZ(0))
+
+        f0 = dev.execute(tape)
+        tapes, fn = finite_diff(tape, approx_order=1, f0=f0)
+
+        # one tape per parameter, plus one global call
+        assert len(tapes) == tape.num_params
+
+    def test_independent_parameters(self):
+        """Test the case where expectation values are independent of some parameters. For those
+        parameters, the gradient should be evaluated to zero without executing the device."""
+        dev = qml.device("default.qubit", wires=2)
+
+        with qml.tape.JacobianTape() as tape1:
+            qml.RX(1, wires=[0])
+            qml.RX(1, wires=[1])
+            qml.expval(qml.PauliZ(0))
+
+        with qml.tape.JacobianTape() as tape2:
+            qml.RX(1, wires=[0])
+            qml.RX(1, wires=[1])
+            qml.expval(qml.PauliZ(1))
+
+        tapes, fn = finite_diff(tape1, approx_order=1)
+        j1 = fn(dev.batch_execute(tapes))
+
+        # We should only be executing the device to differentiate 1 parameter (2 executions)
+        assert dev.num_executions == 2
+
+        tapes, fn = finite_diff(tape2, approx_order=1)
+        j2 = fn(dev.batch_execute(tapes))
+
+        exp = -np.sin(1)
+
+        assert np.allclose(j1, [exp, 0])
+        assert np.allclose(j2, [0, exp])
+
+
+@pytest.mark.parametrize("approx_order", [2, 4])
+@pytest.mark.parametrize("strategy", ["forward", "backward", "center"])
+class TestFiniteDiffIntegration:
+    """Tests for the finite difference gradient transform"""
+
+    def test_ragged_output(self, approx_order, strategy):
+        """Test that the Jacobian is correctly returned for a tape
+        with ragged output"""
+        dev = qml.device("default.qubit", wires=3)
+        params = [1.0, 1.0, 1.0]
+
+        with qml.tape.JacobianTape() as tape:
+            qml.RX(params[0], wires=[0])
+            qml.RY(params[1], wires=[1])
+            qml.RZ(params[2], wires=[2])
+            qml.CNOT(wires=[0, 1])
+            qml.probs(wires=0)
+            qml.probs(wires=[1, 2])
+
+        tapes, fn = finite_diff(tape, approx_order=approx_order, strategy=strategy)
+        res = fn(dev.batch_execute(tapes))
+        assert res.shape == (6, 3)
+
+    def test_single_expectation_value(self, approx_order, strategy, tol):
+        """Tests correct output shape and evaluation for a tape
+        with a single expval output"""
+        dev = qml.device("default.qubit", wires=2)
+        x = 0.543
+        y = -0.654
+
+        with qml.tape.JacobianTape() as tape:
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        tapes, fn = finite_diff(tape, approx_order=approx_order, strategy=strategy)
+        res = fn(dev.batch_execute(tapes))
+        assert res.shape == (1, 2)
+
+        expected = np.array([[-np.sin(y) * np.sin(x), np.cos(y) * np.cos(x)]])
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_single_expectation_value_with_argnum_all(self, approx_order, strategy, tol):
+        """Tests correct output shape and evaluation for a tape
+        with a single expval output where all parameters are chosen to compute
+        the jacobian"""
+        dev = qml.device("default.qubit", wires=2)
+        x = 0.543
+        y = -0.654
+
+        with qml.tape.JacobianTape() as tape:
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        # we choose both trainable parameters
+        tapes, fn = finite_diff(tape, argnum=[0, 1], approx_order=approx_order, strategy=strategy)
+        res = fn(dev.batch_execute(tapes))
+        assert res.shape == (1, 2)
+
+        expected = np.array([[-np.sin(y) * np.sin(x), np.cos(y) * np.cos(x)]])
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_single_expectation_value_with_argnum_one(self, approx_order, strategy, tol):
+        """Tests correct output shape and evaluation for a tape
+        with a single expval output where only one parameter is chosen to
+        estimate the jacobian.
+
+        This test relies on the fact that exactly one term of the estimated
+        jacobian will match the expected analytical value.
+        """
+        dev = qml.device("default.qubit", wires=2)
+        x = 0.543
+        y = -0.654
+
+        with qml.tape.JacobianTape() as tape:
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        # we choose only 1 trainable parameter
+        tapes, fn = finite_diff(tape, argnum=1, approx_order=approx_order, strategy=strategy)
+        res = fn(dev.batch_execute(tapes))
+        assert res.shape == (1, 2)
+
+        expected = np.array([[0, np.cos(y) * np.cos(x)]])
+        res = res.flatten()
+        expected = expected.flatten()
+
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_multiple_expectation_values(self, approx_order, strategy, tol):
+        """Tests correct output shape and evaluation for a tape
+        with multiple expval outputs"""
+        dev = qml.device("default.qubit", wires=2)
+        x = 0.543
+        y = -0.654
+
+        with qml.tape.JacobianTape() as tape:
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0))
+            qml.expval(qml.PauliX(1))
+
+        tapes, fn = finite_diff(tape, approx_order=approx_order, strategy=strategy)
+        res = fn(dev.batch_execute(tapes))
+        assert res.shape == (2, 2)
+
+        expected = np.array([[-np.sin(x), 0], [0, np.cos(y)]])
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_var_expectation_values(self, approx_order, strategy, tol):
+        """Tests correct output shape and evaluation for a tape
+        with expval and var outputs"""
+        dev = qml.device("default.qubit", wires=2)
+        x = 0.543
+        y = -0.654
+
+        with qml.tape.JacobianTape() as tape:
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0))
+            qml.var(qml.PauliX(1))
+
+        tapes, fn = finite_diff(tape, approx_order=approx_order, strategy=strategy)
+        res = fn(dev.batch_execute(tapes))
+        assert res.shape == (2, 2)
+
+        expected = np.array([[-np.sin(x), 0], [0, -2 * np.cos(y) * np.sin(y)]])
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_prob_expectation_values(self, approx_order, strategy, tol):
+        """Tests correct output shape and evaluation for a tape
+        with prob and expval outputs"""
+        dev = qml.device("default.qubit", wires=2)
+        x = 0.543
+        y = -0.654
+
+        with qml.tape.JacobianTape() as tape:
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0))
+            qml.probs(wires=[0, 1])
+
+        tapes, fn = finite_diff(tape, approx_order=approx_order, strategy=strategy)
+        res = fn(dev.batch_execute(tapes))
+        assert res.shape == (5, 2)
+
+        expected = (
+            np.array(
+                [
+                    [-2 * np.sin(x), 0],
+                    [
+                        -(np.cos(y / 2) ** 2 * np.sin(x)),
+                        -(np.cos(x / 2) ** 2 * np.sin(y)),
+                    ],
+                    [
+                        -(np.sin(x) * np.sin(y / 2) ** 2),
+                        (np.cos(x / 2) ** 2 * np.sin(y)),
+                    ],
+                    [
+                        (np.sin(x) * np.sin(y / 2) ** 2),
+                        (np.sin(x / 2) ** 2 * np.sin(y)),
+                    ],
+                    [
+                        (np.cos(y / 2) ** 2 * np.sin(x)),
+                        -(np.sin(x / 2) ** 2 * np.sin(y)),
+                    ],
+                ]
+            )
+            / 2
+        )
+
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+
+@pytest.mark.parametrize("approx_order", [2])
+@pytest.mark.parametrize("strategy", ["center"])
+class TestFiniteDiffGradients:
+    """Test that the transform is differentiable"""
+
+    def test_autograd(self, approx_order, strategy, tol):
+        """Tests that the output of the finite-difference transform
+        can be differentiated using autograd, yielding second derivatives."""
+        dev = qml.device("default.qubit.autograd", wires=2)
+        params = np.array([0.543, -0.654], requires_grad=True)
+
+        def cost_fn(x):
+            with qml.tape.JacobianTape() as tape:
+                qml.RX(x[0], wires=[0])
+                qml.RY(x[1], wires=[1])
+                qml.CNOT(wires=[0, 1])
+                qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+            tape.trainable_params = {0, 1}
+            tapes, fn = finite_diff(tape, n=1, approx_order=approx_order, strategy=strategy)
+            jac = fn(dev.batch_execute(tapes))
+            return jac
+
+        res = qml.jacobian(cost_fn)(params)
+        x, y = params
+        expected = np.array(
+            [
+                [-np.cos(x) * np.sin(y), -np.cos(y) * np.sin(x)],
+                [-np.cos(y) * np.sin(x), -np.cos(x) * np.sin(y)],
+            ]
+        )
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_autograd_ragged(self, approx_order, strategy, tol):
+        """Tests that the output of the finite-difference transform
+        of a ragged tape can be differentiated using autograd, yielding second derivatives."""
+        dev = qml.device("default.qubit.autograd", wires=2)
+        params = np.array([0.543, -0.654], requires_grad=True)
+
+        def cost_fn(x):
+            with qml.tape.JacobianTape() as tape:
+                qml.RX(x[0], wires=[0])
+                qml.RY(x[1], wires=[1])
+                qml.CNOT(wires=[0, 1])
+                qml.expval(qml.PauliZ(0))
+                qml.probs(wires=[1])
+
+            tape.trainable_params = {0, 1}
+            tapes, fn = finite_diff(tape, n=1, approx_order=approx_order, strategy=strategy)
+            jac = fn(dev.batch_execute(tapes))
+            return jac[1, 0]
+
+        x, y = params
+        res = qml.grad(cost_fn)(params)
+        expected = np.array([-np.cos(x) * np.cos(y) / 2, np.sin(x) * np.sin(y) / 2])
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_tf(self, approx_order, strategy, tol):
+        """Tests that the output of the finite-difference transform
+        can be differentiated using TF, yielding second derivatives."""
+        tf = pytest.importorskip("tensorflow")
+
+        dev = qml.device("default.qubit.tf", wires=2)
+        params = tf.Variable([0.543, -0.654], dtype=tf.float64)
+
+        with tf.GradientTape() as t:
+            with qml.tape.JacobianTape() as tape:
+                qml.RX(params[0], wires=[0])
+                qml.RY(params[1], wires=[1])
+                qml.CNOT(wires=[0, 1])
+                qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+            tape.trainable_params = {0, 1}
+            tapes, fn = finite_diff(tape, n=1, approx_order=approx_order, strategy=strategy)
+            jac = fn(dev.batch_execute(tapes))
+
+        x, y = 1.0 * params
+
+        expected = np.array([-np.sin(x) * np.sin(y), np.cos(x) * np.cos(y)])
+        assert np.allclose(jac, expected, atol=tol, rtol=0)
+
+        res = t.jacobian(jac, params)
+        expected = np.array(
+            [
+                [-np.cos(x) * np.sin(y), -np.cos(y) * np.sin(x)],
+                [-np.cos(y) * np.sin(x), -np.cos(x) * np.sin(y)],
+            ]
+        )
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_tf_ragged(self, approx_order, strategy, tol):
+        """Tests that the output of the finite-difference transform
+        of a ragged tape can be differentiated using TF, yielding second derivatives."""
+        tf = pytest.importorskip("tensorflow")
+        dev = qml.device("default.qubit.tf", wires=2)
+        params = tf.Variable([0.543, -0.654], dtype=tf.float64)
+
+        with tf.GradientTape() as t:
+            with qml.tape.JacobianTape() as tape:
+                qml.RX(params[0], wires=[0])
+                qml.RY(params[1], wires=[1])
+                qml.CNOT(wires=[0, 1])
+                qml.expval(qml.PauliZ(0))
+                qml.probs(wires=[1])
+
+            tape.trainable_params = {0, 1}
+            tapes, fn = finite_diff(tape, n=1, approx_order=approx_order, strategy=strategy)
+            jac = fn(dev.batch_execute(tapes))[1, 0]
+
+        x, y = 1.0 * params
+        res = t.gradient(jac, params)
+        expected = np.array([-np.cos(x) * np.cos(y) / 2, np.sin(x) * np.sin(y) / 2])
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_torch(self, approx_order, strategy, tol):
+        """Tests that the output of the finite-difference transform
+        can be differentiated using Torch, yielding second derivatives."""
+        torch = pytest.importorskip("torch")
+        from pennylane.interfaces.torch import TorchInterface
+
+        dev = qml.device("default.qubit", wires=2)
+        params = torch.tensor([0.543, -0.654], dtype=torch.float64, requires_grad=True)
+
+        with TorchInterface.apply(qml.tape.QubitParamShiftTape()) as tape:
+            qml.RX(params[0], wires=[0])
+            qml.RY(params[1], wires=[1])
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        tapes, fn = finite_diff(tape, n=1, approx_order=approx_order, strategy=strategy)
+        jac = fn([t.execute(dev) for t in tapes])
+        cost = torch.sum(jac)
+        cost.backward()
+        hess = params.grad
+
+        x, y = params.detach().numpy()
+
+        expected = np.array([-np.sin(x) * np.sin(y), np.cos(x) * np.cos(y)])
+        assert np.allclose(jac.detach().numpy(), expected, atol=tol, rtol=0)
+
+        expected = np.array(
+            [
+                [-np.cos(x) * np.sin(y), -np.cos(y) * np.sin(x)],
+                [-np.cos(y) * np.sin(x), -np.cos(x) * np.sin(y)],
+            ]
+        )
+        assert np.allclose(hess.detach().numpy(), np.sum(expected, axis=0), atol=tol, rtol=0)
+
+    def test_jax(self, approx_order, strategy, tol):
+        """Tests that the output of the finite-difference transform
+        can be differentiated using JAX, yielding second derivatives."""
+        jax = pytest.importorskip("jax")
+        from jax import numpy as jnp
+        from pennylane.interfaces.jax import JAXInterface
+        from jax.config import config
+
+        config.update("jax_enable_x64", True)
+
+        dev = qml.device("default.qubit", wires=2)
+        params = jnp.array([0.543, -0.654])
+
+        def cost_fn(x):
+            with JAXInterface.apply(qml.tape.QubitParamShiftTape()) as tape:
+                qml.RX(x[0], wires=[0])
+                qml.RY(x[1], wires=[1])
+                qml.CNOT(wires=[0, 1])
+                qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+            tape.trainable_params = {0, 1}
+            tapes, fn = finite_diff(tape, n=1, approx_order=approx_order, strategy=strategy)
+            jac = fn([t.execute(dev) for t in tapes])
+            return jac
+
+        res = jax.jacobian(cost_fn)(params)
+        x, y = params
+        expected = np.array(
+            [
+                [-np.cos(x) * np.sin(y), -np.cos(y) * np.sin(x)],
+                [-np.cos(y) * np.sin(x), -np.cos(x) * np.sin(y)],
+            ]
+        )
+        assert np.allclose(res, expected, atol=tol, rtol=0)
diff --git a/tests/interfaces/test_qnode_jax.py b/tests/interfaces/test_qnode_jax.py
index 0e48da3b4..a1be16d19 100644
--- a/tests/interfaces/test_qnode_jax.py
+++ b/tests/interfaces/test_qnode_jax.py
@@ -53,7 +53,7 @@ def test_to_jax():
     assert "DeviceArray" in val.__repr__()
 
 
-def test_simple_jacobian():
+def test_simple_jacobian(tol):
     """Test the use of jax.jaxrev"""
     dev = qml.device("default.mixed", wires=2)  # A non-JAX device.
 
@@ -69,7 +69,7 @@ def test_simple_jacobian():
     # of a numpy array.
     assert "DeviceArray" in grads.__repr__()
     assert grads.shape == (2,)
-    np.testing.assert_allclose(grads, np.array([-0.09784342, -0.19767685]))
+    np.testing.assert_allclose(grads, np.array([-0.09784342, -0.19767685]), atol=tol, rtol=0)
 
 
 def test_simple_grad():