Repository: PennyLaneAI/pennylane
Pull request #1501 (2021-08-05 17:59:35)

Differentiable batch execute using autograd
**Context:** This PR adds support for differentiable batch execution of circuits using autograd.

**Description of the change:**

A new subpackage has been added, `pennylane/interfaces/batch`. This is a temporary directory --- once `batch_execute` is the default pathway in PL, the old interfaces will be removed.

This PR adds the following:

- `qml.interfaces.batch.execute()` - a high level wrapper function, that dispatches to the correct autodiff interface. Currently, only autograd is supported. In addition, this high-level wrapper also contains shared code for determining the correct forward/backward execution function.

- `qml.interfaces.autograd` - a module containing an autograd custom gradient function for `batch_execute`.

- `Device.execute_and_gradients`: new Device method for computing tape results alongside tape gradients

- `Device.gradients`: new Device method for computing tape gradients. Can be considered the gradient equivalent of `Device.batch_execute`.

**Benefits:**

- Execute tapes in a batch, with the output remaining differentiable.

- Supports gradient transforms and device execution methods

- Since gradient transforms are differentiable, nth order higher derivatives are supported. Compared to PL master, this allows nth order derivatives of _everything_, including expval, var, probs, tensor products, non-two-term shifts, etc.

- Supports both forward and backward accumulation modes.

  - Gradient transforms only support `mode="backward"`, due to there not being much of a use-case for computing shift-rule gradients during the forward pass.

  - Device methods support `mode="forward"` and `mode="backward"`. Forward mode is the default when using a device method, since it is often the case that the device state on the forward pass is re-used to compute gradients.

**Potential drawbacks:**

- The recursive method used to compute higher-derivatives is not so smart. For example, when computing the Hessian, all matrix elements are computed, rather than just the upper-tridiagonal.

- Device gradient methods are not differentiable, and therefore when using device methods, higher order derivatives are not accessible unless directly supported by the device method.

- In #1493, work is being done to create a standardized API for gradient transforms. Until then, we simply assume that any `gradient_fn` within the `pennylane.gradients` module is a transform.

- All gradient transforms, and all device gradients (e.g., adjoint) are supported. The reversible method, however, is not currently supported, since it is not a transform nor a device method.

**Issues:** n/a

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [.github/CHANGELOG.md](/PennyLaneAI/pennylane/blob/master/.github/CHANGELOG.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1501?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1501](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1501?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (89bdd8d) into [vjp-transform](https://codecov.io/gh/PennyLaneAI/pennylane/commit/35e1848ba0731c5d2ccc68c6aa604b580964aa06?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (35e1848) will **decrease** coverage by `9.44%`.
> The diff coverage is `7.35%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1501/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/1501?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@                Coverage Diff                @@
##           vjp-transform    #1501      +/-   ##
=================================================
- Coverage          98.37%   88.92%   -9.45%     
=================================================
  Files                183      180       -3     
  Lines              12978    12514     -464     
=================================================
- Hits               12767    11128    -1639     
- Misses               211     1386    +1175     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1501?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/interfaces/batch/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1501/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvX19pbml0X18ucHk=) | `0.00% <0.00%> (ø)` | |
| [pennylane/interfaces/batch/autograd.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1501/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvYXV0b2dyYWQucHk=) | `0.00% <0.00%> (ø)` | |
| [pennylane/gradients/vjp.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1501/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy92anAucHk=) | `79.24% <15.38%> (-20.76%)` | :arrow_down: |
| [pennylane/\_qubit\_device.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1501/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19xdWJpdF9kZXZpY2UucHk=) | `97.67% <22.22%> (-0.98%)` | :arrow_down: |
| [pennylane/gradients/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1501/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/beta/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1501/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2JldGEvX19pbml0X18ucHk=) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/devices/tf\_ops.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1501/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvdGZfb3BzLnB5) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/devices/jax\_ops.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1501/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvamF4X29wcy5weQ==) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/interfaces/torch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1501/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvdG9yY2gucHk=) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/beta/devices/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1501/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2JldGEvZGV2aWNlcy9fX2luaXRfXy5weQ==) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| ... and [56 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1501/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1501?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/1501?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [35e1848...89bdd8d](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1501?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: @co9olguy: have taken into account your suggestions, and written tests. This is now ready for proper review 🙂 

Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index 0c5204cea..243b8afa2 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -52,6 +52,32 @@
   (array([-0.12777055,  0.0166009 ]), array(0.0917819))
   ```
 
+* Support for differentiable execution of batches of circuits has been
+  added, via the beta `pennylane.batch` module.
+  [(#1501)](https://github.com/PennyLaneAI/pennylane/pull/1501)
+
+  For example:
+
+  ```python
+  def cost_fn(x):
+      with qml.tape.JacobianTape() as tape1:
+          qml.RX(x[0], wires=[0])
+          qml.RY(x[1], wires=[1])
+          qml.CNOT(wires=[0, 1])
+          qml.var(qml.PauliZ(0) @ qml.PauliX(1))
+
+      with qml.tape.JacobianTape() as tape2:
+          qml.RX(x[0], wires=0)
+          qml.RY(x[0], wires=1)
+          qml.CNOT(wires=[0, 1])
+          qml.probs(wires=1)
+
+      result = execute([tape1, tape2], dev, gradient_fn=param_shift)
+      return result[0] + result[1][0, 0]
+
+  res = qml.grad(cost_fn)(params)
+  ```
+
 <h3>Improvements</h3>
 
 * The `group_observables` transform is now differentiable.
diff --git a/pennylane/_device.py b/pennylane/_device.py
index 637c79cc6..32e165264 100644
--- a/pennylane/_device.py
+++ b/pennylane/_device.py
@@ -496,6 +496,64 @@ class Device(abc.ABC):
 
         return results
 
+    def execute_and_gradients(self, circuits, method="jacobian", **kwargs):
+        """Execute a batch of quantum circuits on the device, and return both the
+        results and the gradients.
+
+        The circuits are represented by tapes, and they are executed
+        one-by-one using the device's ``execute`` method. The results and the
+        corresponding Jacobians are collected in a list.
+
+        For plugin developers: This method should be overwritten if the device
+        can efficiently run multiple circuits on a backend, for example using
+        parallel and/or asynchronous executions, and return both the results and the
+        Jacobians.
+
+        Args:
+            circuits (list[.tape.QuantumTape]): circuits to execute on the device
+            method (str): the device method to call to compute the Jacobian of a single circuit
+            **kwargs: keyword argument to pass when calling ``method``
+
+        Returns:
+            tuple[list[array[float]], list[array[float]]]: Tuple containing list of measured value(s)
+            and list of Jacobians. Returned Jacobians should be of shape ``(output_shape, num_params)``.
+        """
+        gradient_method = getattr(self, method)
+
+        res = []
+        jacs = []
+
+        for circuit in circuits:
+            # Evaluations and gradients are paired, so that
+            # devices can re-use the device state for the
+            # gradient computation (if applicable).
+            res.append(circuit.execute(self))
+            jacs.append(gradient_method(circuit, **kwargs))
+
+        return res, jacs
+
+    def gradients(self, circuits, method="jacobian", **kwargs):
+        """Return the gradients of a batch of quantum circuits on the device.
+
+        The gradient method ``method`` is called sequentially for each
+        circuit, and the corresponding Jacobians are collected in a list.
+
+        For plugin developers: This method should be overwritten if the device
+        can efficiently compute the gradient of multiple circuits on a
+        backend, for example using parallel and/or asynchronous executions.
+
+        Args:
+            circuits (list[.tape.QuantumTape]): circuits to execute on the device
+            method (str): the device method to call to compute the Jacobian of a single circuit
+            **kwargs: keyword argument to pass when calling ``method``
+
+        Returns:
+            list[array[float]]: List of Jacobians. Returned Jacobians should be of
+            shape ``(output_shape, num_params)``.
+        """
+        gradient_method = getattr(self, method)
+        return [gradient_method(circuit, **kwargs) for circuit in circuits]
+
     @property
     def op_queue(self):
         """The operation queue to be applied.
diff --git a/pennylane/gradients/vjp.py b/pennylane/gradients/vjp.py
index 7e68eb3bd..eb2baed5b 100644
--- a/pennylane/gradients/vjp.py
+++ b/pennylane/gradients/vjp.py
@@ -48,9 +48,9 @@ def compute_vjp(dy, jac):
     return math.tensordot(jac, dy_row, [[0], [0]])
 
 
-def vjp(tape, dy, gradient_fn):
+def vjp(tape, dy, gradient_fn, gradient_kwargs=None):
     r"""Generate the gradient tapes and processing function required to compute
-    the vector-Jacobian product of a tape.
+    the vector-Jacobian products of a tape.
 
     Consider a function :math:`\mathbf{f}(\mathbf{x})`. The Jacobian is given by
 
@@ -88,6 +88,8 @@ def vjp(tape, dy, gradient_fn):
             matching the output shape of the corresponding tape.
         gradient_fn (callable): the gradient transform to use to differentiate
             the tape
+        gradient_kwargs (dict): dictionary of keyword arguments to pass when
+            determining the gradients of tapes
 
     Returns:
         tensor_like or None: Vector-Jacobian product. Returns None if the tape
@@ -143,6 +145,7 @@ def vjp(tape, dy, gradient_fn):
     """
     # t._par_info = {}
     # t._update()
+    gradient_kwargs = gradient_kwargs or {}
     num_params = len(tape.trainable_params)
 
     if num_params == 0:
@@ -156,7 +159,7 @@ def vjp(tape, dy, gradient_fn):
         # and we can avoid a quantum computation.
         return [], lambda _: math.convert_like(np.zeros([num_params]), dy)
 
-    gradient_tapes, fn = gradient_fn(tape)
+    gradient_tapes, fn = gradient_fn(tape, **gradient_kwargs)
 
     def processing_fn(results):
         # postprocess results to compute the Jacobian
@@ -166,7 +169,7 @@ def vjp(tape, dy, gradient_fn):
     return gradient_tapes, processing_fn
 
 
-def batch_vjp(tapes, dys, gradient_fn, reduction="append"):
+def batch_vjp(tapes, dys, gradient_fn, reduction="append", gradient_kwargs=None):
     r"""Generate the gradient tapes and processing function required to compute
     the vector-Jacobian products of a batch of tapes.
 
@@ -211,6 +214,8 @@ def batch_vjp(tapes, dys, gradient_fn, reduction="append"):
             If ``append``, then the output of the function will be of the form
             ``List[tensor_like]``, with each element corresponding to the VJP of each
             input tape. If ``extend``, then the output VJPs will be concatenated.
+        gradient_kwargs (dict): dictionary of keyword arguments to pass when
+            determining the gradients of tapes
 
     Returns:
         List[tensor_like or None]: list of vector-Jacobian products. ``None`` elements corresponds
@@ -278,13 +283,14 @@ def batch_vjp(tapes, dys, gradient_fn, reduction="append"):
     tensor([[-4.7924e-01, -9.0857e-01, -2.4198e-01],
             [-9.2973e-02, -1.0772e+00,  4.7184e-09]], dtype=torch.float64)
     """
+    gradient_kwargs = gradient_kwargs or {}
     reshape_info = []
     gradient_tapes = []
     processing_fns = []
 
     # Loop through the tapes and dys vector
     for tape, dy in zip(tapes, dys):
-        g_tapes, fn = vjp(tape, dy, gradient_fn)
+        g_tapes, fn = vjp(tape, dy, gradient_fn, gradient_kwargs)
 
         reshape_info.append(len(g_tapes))
         processing_fns.append(fn)
diff --git a/pennylane/interfaces/batch/__init__.py b/pennylane/interfaces/batch/__init__.py
new file mode 100644
index 000000000..2dd08bb18
--- /dev/null
+++ b/pennylane/interfaces/batch/__init__.py
@@ -0,0 +1,131 @@
+# 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 defines functions for interfacing devices' batch execution
+capabilities with different machine learning libraries.
+"""
+# pylint: disable=import-outside-toplevel,too-many-arguments
+import pennylane as qml
+
+from .autograd import execute as execute_autograd
+
+
+def execute(tapes, device, gradient_fn, interface="autograd", mode="best", gradient_kwargs=None):
+    """Execute a batch of tapes on a device in an autodifferentiable-compatible manner.
+
+    Args:
+        tapes (Sequence[.QuantumTape]): batch of tapes to execute
+        device (.Device): Device to use to execute the batch of tapes.
+            If the device does not provide a ``batch_execute`` method,
+            by default the tapes will be executed in serial.
+        gradient_fn (None or callable): The gradient transform function to use
+            for backward passes. If "device", the device will be queried directly
+            for the gradient (if supported).
+        interface (str): The interface that will be used for classical autodifferentiation.
+            This affects the types of parameters that can exist on the input tapes.
+            Available options include ``autograd``, ``torch``, ``tf``, and ``jax``.
+        mode (str): Whether the gradients should be computed on the forward
+            pass (``forward``) or the backward pass (``backward``). Only applies
+            if the device is queried for the gradient; gradient transform
+            functions available in ``qml.gradients`` are only supported on the backward
+            pass.
+        gradient_kwargs (dict): dictionary of keyword arguments to pass when
+            determining the gradients of tapes
+
+    Returns:
+        list[list[float]]: A nested list of tape results. Each element in
+        the returned list corresponds in order to the provided tapes.
+
+    **Example**
+
+    Consider the following cost function:
+
+    .. code-block:: python
+
+        dev = qml.device("lightning.qubit", wires=2)
+
+        def cost_fn(params, x):
+            with qml.tape.JacobianTape() as tape1:
+                qml.RX(params[0], wires=0)
+                qml.RY(params[1], wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            with qml.tape.JacobianTape() as tape2:
+                qml.RX(params[2], wires=0)
+                qml.RY(x[0], wires=1)
+                qml.CNOT(wires=[0, 1])
+                qml.probs(wires=0)
+
+            tapes = [tape1, tape2]
+
+            # execute both tapes in a batch on the given device
+            res = execute(tapes, dev, qml.gradients.param_shift)
+
+            return res[0][0] + res[1][0, 0] - res[1][0, 1]
+
+    In this cost function, two **independent** quantum tapes are being
+    constructed; one returning an expectation value, the other probabilities.
+    We then batch execute the two tapes, and reduce the results to obtain
+    a scalar.
+
+    Let's execute this cost function while tracking the gradient:
+
+    >>> params = np.array([0.1, 0.2, 0.3], requires_grad=True)
+    >>> x = np.array([0.5], requires_grad=True)
+    >>> cost_fn(params, x)
+    1.9305068163274222
+
+    Since the ``execute`` function is differentiable, we can
+    also compute the gradient:
+
+    >>> qml.grad(cost_fn)(params, x)
+    (array([-0.0978434 , -0.19767681, -0.29552021]), array([5.37764278e-17]))
+
+    Finally, we can also compute any nth-order derivative. Let's compute the Jacobian
+    of the gradient (that is, the Hessian):
+
+    >>> x.requires_grad = False
+    >>> qml.jacobian(qml.grad(cost_fn))(params, x)
+    array([[-0.97517033,  0.01983384,  0.        ],
+           [ 0.01983384, -0.97517033,  0.        ],
+           [ 0.        ,  0.        , -0.95533649]])
+    """
+    # Default execution function; simply call device.batch_execute
+    # and return no Jacobians.
+    execute_fn = lambda tapes, **kwargs: (device.batch_execute(tapes), [])
+    gradient_kwargs = gradient_kwargs or {}
+
+    if gradient_fn == "device":
+        # gradient function is a device method
+
+        if mode in ("forward", "best"):
+            # replace the forward execution function to return
+            # both results and gradients
+            execute_fn = device.execute_and_gradients
+            gradient_fn = None
+
+        elif mode == "backward":
+            # replace the backward gradient computation
+            gradient_fn = device.gradients
+
+    elif mode == "forward":
+        # In "forward" mode, gradients are automatically handled
+        # within execute_and_gradients, so providing a gradient_fn
+        # in this case would have ambiguous behaviour.
+        raise ValueError("Gradient transforms cannot be used with mode='forward'")
+
+    if interface == "autograd":
+        return execute_autograd(tapes, device, execute_fn, gradient_fn, gradient_kwargs)
+
+    raise ValueError(f"Unknown interface {interface}")
diff --git a/pennylane/interfaces/batch/autograd.py b/pennylane/interfaces/batch/autograd.py
new file mode 100644
index 000000000..252caafed
--- /dev/null
+++ b/pennylane/interfaces/batch/autograd.py
@@ -0,0 +1,205 @@
+# 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 adding the Autograd interface
+to a PennyLane Device class.
+"""
+# pylint: disable=too-many-arguments
+import inspect
+
+import autograd
+from autograd.numpy.numpy_boxes import ArrayBox
+
+import pennylane as qml
+from pennylane import numpy as np
+
+
+def execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1):
+    """Execute a batch of tapes with Autograd parameters on a device.
+
+    Args:
+        tapes (Sequence[.QuantumTape]): batch of tapes to execute
+        device (.Device): Device to use to execute the batch of tapes.
+            If the device does not provide a ``batch_execute`` method,
+            by default the tapes will be executed in serial.
+        execute_fn (callable): The execution function used to execute the tapes
+            during the forward pass. This function must return a tuple ``(results, jacobians)``.
+            If ``jacobians`` is an empty list, then ``gradient_fn`` is used to
+            compute the gradients during the backwards pass.
+        gradient_kwargs (dict): dictionary of keyword arguments to pass when
+            determining the gradients of tapes
+        gradient_fn (callable): the gradient function to use to compute quantum gradients
+        _n (int): a positive integer used to track nesting of derivatives, for example
+            if the nth-order derivative is requested.
+
+    Returns:
+        list[list[float]]: A nested list of tape results. Each element in
+        the returned list corresponds in order to the provided tapes.
+    """
+    for tape in tapes:
+        # set the trainable parameters
+        params = tape.get_parameters(trainable_only=False)
+        tape.trainable_params = qml.math.get_trainable_indices(params)
+
+    parameters = autograd.builtins.tuple(
+        [autograd.builtins.list(t.get_parameters()) for t in tapes]
+    )
+
+    return _execute(
+        parameters,
+        tapes=tapes,
+        device=device,
+        execute_fn=execute_fn,
+        gradient_fn=gradient_fn,
+        gradient_kwargs=gradient_kwargs,
+        _n=_n,
+    )[0]
+
+
+@autograd.extend.primitive
+def _execute(
+    parameters,
+    tapes=None,
+    device=None,
+    execute_fn=None,
+    gradient_fn=None,
+    gradient_kwargs=None,
+    _n=1,
+):  # pylint: disable=dangerous-default-value,unused-argument
+    """Autodifferentiable wrapper around ``Device.batch_execute``.
+
+    The signature of this function is designed to work around Autograd restrictions.
+    Note that the ``parameters`` argument is dependent on the ``tapes`` argument;
+    this function should always be called as follows:
+
+    >>> parameters = [autograd.builtins.list(t.get_parameters()) for t in tapes])
+    >>> parameters = autograd.builtins.tuple(parameters)
+    >>> _batch_execute(parameters, tapes=tapes, device=device)
+
+    In particular:
+
+    - ``parameters`` is dependent on the provided tapes: always extract them as above
+    - ``tapes`` is a *required* argument
+    - ``device`` is a *required* argument
+
+    The private argument ``_n`` is used to track nesting of derivatives, for example
+    if the nth-order derivative is requested. Do not set this argument unless you
+    understand the consequences!
+    """
+    with qml.tape.Unwrap(*tapes):
+        res, jacs = execute_fn(tapes, **gradient_kwargs)
+
+    for i, r in enumerate(res):
+        res[i] = np.tensor(r)
+
+        if r.dtype == np.dtype("object"):
+            # For backwards compatibility, we flatten ragged tape outputs
+            res[i] = np.hstack(r)
+
+    return res, jacs
+
+
+def vjp(
+    ans,
+    parameters,
+    tapes=None,
+    device=None,
+    execute_fn=None,
+    gradient_fn=None,
+    gradient_kwargs=None,
+    _n=1,
+):  # pylint: disable=dangerous-default-value,unused-argument
+    """Returns the vector-Jacobian product operator for a batch of quantum tapes.
+
+    Args:
+        ans (array): the result of the batch tape execution
+        parameters (list[list[Any]]): Nested list of the quantum tape parameters.
+            This argument should be generated from the provided list of tapes.
+        tapes (Sequence[.QuantumTape]): batch of tapes to execute
+        device (.Device): Device to use to execute the batch of tapes.
+            If the device does not provide a ``batch_execute`` method,
+            by default the tapes will be executed in serial.
+        execute_fn (callable): The execution function used to execute the tapes
+            during the forward pass. This function must return a tuple ``(results, jacobians)``.
+            If ``jacobians`` is an empty list, then ``gradient_fn`` is used to
+            compute the gradients during the backwards pass.
+        gradient_fn (callable): the gradient function to use to compute quantum gradients
+        gradient_kwargs (dict): dictionary of keyword arguments to pass when
+            determining the gradients of tapes
+        _n (int): a positive integer used to track nesting of derivatives, for example
+            if the nth-order derivative is requested.
+
+    Returns:
+        function: this function accepts the backpropagation
+        gradient output vector, and computes the vector-Jacobian product
+    """
+
+    def grad_fn(dy):
+        """Returns the vector-Jacobian product with given
+        parameter values and output gradient dy"""
+
+        dy = dy[0]
+        jacs = ans[1]
+
+        if jacs:
+            # Jacobians were computed on the forward pass (mode="forward")
+            # No additional quantum evaluations needed; simply compute the VJPs directly.
+            vjps = [qml.gradients.compute_vjp(d, jac) for d, jac in zip(dy, jacs)]
+
+        else:
+            # Need to compute the Jacobians on the backward pass (accumulation="backward")
+
+            # Temporary: check if the gradient function is a differentiable transform.
+            # For the moment, simply check if it is part of the `qml.gradients` package.
+            # Longer term, we should have a way of checking this directly
+            # (e.g., isinstance(gradient_fn, GradientTransform))
+            module_name = getattr(inspect.getmodule(gradient_fn), "__name__", "")
+
+            if "pennylane.gradients" in module_name:
+
+                # Generate and execute the required gradient tapes
+                vjp_tapes, processing_fn = qml.gradients.batch_vjp(
+                    tapes, dy, gradient_fn, reduction="append", gradient_kwargs=gradient_kwargs
+                )
+
+                # This is where the magic happens. Note that we call ``execute``.
+                # This recursion, coupled with the fact that the gradient transforms
+                # are differentiable, allows for arbitrary order differentiation.
+                vjps = processing_fn(
+                    execute(vjp_tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=_n + 1)
+                )
+
+            elif inspect.ismethod(gradient_fn) and gradient_fn.__self__ is device:
+                # Gradient function is a device method.
+                # Note that unlike the previous branch:
+                #
+                # - there is no recursion here
+                # - gradient_fn is not differentiable
+                #
+                # so we cannot support higher-order derivatives.
+
+                with qml.tape.Unwrap(*tapes):
+                    jacs = gradient_fn(tapes, **gradient_kwargs)
+
+                vjps = [qml.gradients.compute_vjp(d, jac) for d, jac in zip(dy, jacs)]
+
+            else:
+                raise ValueError("Unknown gradient function.")
+
+        return [qml.math.to_numpy(v, max_depth=_n) if isinstance(v, ArrayBox) else v for v in vjps]
+
+    return grad_fn
+
+
+autograd.extend.defvjp(_execute, vjp, argnums=[0])


Test patch:
diff --git a/tests/gradients/test_vjp.py b/tests/gradients/test_vjp.py
index 4b1627e85..92f4a7b47 100644
--- a/tests/gradients/test_vjp.py
+++ b/tests/gradients/test_vjp.py
@@ -400,7 +400,6 @@ class TestBatchVJP:
 
     def test_reduction_append(self):
         """Test the 'append' reduction strategy"""
-
         dev = qml.device("default.qubit", wires=2)
 
         with qml.tape.JacobianTape() as tape1:
diff --git a/tests/interfaces/test_batch_autograd.py b/tests/interfaces/test_batch_autograd.py
new file mode 100644
index 000000000..435c5004c
--- /dev/null
+++ b/tests/interfaces/test_batch_autograd.py
@@ -0,0 +1,596 @@
+# 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.
+"""Unit tests for the autograd interface"""
+import functools
+
+import pytest
+from pennylane import numpy as np
+
+import pennylane as qml
+from pennylane.gradients import param_shift
+from pennylane.interfaces.batch import execute
+
+
+class TestAutogradExecuteUnitTests:
+    """Unit tests for autograd execution"""
+
+    def test_jacobian_options(self, mocker, tol):
+        """Test setting jacobian options"""
+        spy = mocker.spy(qml.gradients, "param_shift")
+
+        a = np.array([0.1, 0.2], requires_grad=True)
+
+        dev = qml.device("default.qubit", wires=1)
+
+        def cost(a, device):
+            with qml.tape.JacobianTape() as tape:
+                qml.RY(a[0], wires=0)
+                qml.RX(a[1], wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            return execute(
+                [tape], device, gradient_fn=param_shift, gradient_kwargs={"shift": np.pi / 4}
+            )[0]
+
+        res = qml.jacobian(cost)(a, device=dev)
+
+        for args in spy.call_args_list:
+            assert args[1]["shift"] == np.pi / 4
+
+    def test_unknown_gradient_fn_error(self):
+        """Test that an error is raised if an unknown gradient function
+        is passed"""
+        a = np.array([0.1, 0.2], requires_grad=True)
+
+        dev = qml.device("default.qubit", wires=1)
+
+        def cost(a, device):
+            with qml.tape.JacobianTape() as tape:
+                qml.RY(a[0], wires=0)
+                qml.RX(a[1], wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            return execute([tape], device, gradient_fn=lambda x: x)[0]
+
+        with pytest.raises(ValueError, match="Unknown gradient function"):
+            res = qml.jacobian(cost)(a, device=dev)
+
+    def test_incorrect_mode(self):
+        """Test that an error is raised if an gradient transform
+        is used with mode=forward"""
+        a = np.array([0.1, 0.2], requires_grad=True)
+
+        dev = qml.device("default.qubit", wires=1)
+
+        def cost(a, device):
+            with qml.tape.JacobianTape() as tape:
+                qml.RY(a[0], wires=0)
+                qml.RX(a[1], wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            return execute([tape], device, gradient_fn=param_shift, mode="forward")[0]
+
+        with pytest.raises(
+            ValueError, match="Gradient transforms cannot be used with mode='forward'"
+        ):
+            res = qml.jacobian(cost)(a, device=dev)
+
+    def test_unknown_interface(self):
+        """Test that an error is raised if the interface is unknown"""
+        a = np.array([0.1, 0.2], requires_grad=True)
+
+        dev = qml.device("default.qubit", wires=1)
+
+        def cost(a, device):
+            with qml.tape.JacobianTape() as tape:
+                qml.RY(a[0], wires=0)
+                qml.RX(a[1], wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            return execute([tape], device, gradient_fn=param_shift, interface=None)[0]
+
+        with pytest.raises(ValueError, match="Unknown interface"):
+            cost(a, device=dev)
+
+    def test_forward_mode(self, mocker):
+        """Test that forward mode uses the `device.execute_and_gradients` pathway"""
+        dev = qml.device("default.qubit", wires=1)
+        spy = mocker.spy(dev, "execute_and_gradients")
+
+        def cost(a):
+            with qml.tape.JacobianTape() as tape:
+                qml.RY(a[0], wires=0)
+                qml.RX(a[1], wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            return execute(
+                [tape],
+                dev,
+                gradient_fn="device",
+                gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True},
+            )[0]
+
+        a = np.array([0.1, 0.2], requires_grad=True)
+        cost(a)
+
+        # adjoint method only performs a single device execution, but gets both result and gradient
+        assert dev.num_executions == 1
+        spy.assert_called()
+
+    def test_backward_mode(self, mocker):
+        """Test that backward mode uses the `device.batch_execute` and `device.gradients` pathway"""
+        dev = qml.device("default.qubit", wires=1)
+        spy_execute = mocker.spy(qml.devices.DefaultQubit, "batch_execute")
+        spy_gradients = mocker.spy(qml.devices.DefaultQubit, "gradients")
+
+        def cost(a):
+            with qml.tape.JacobianTape() as tape:
+                qml.RY(a[0], wires=0)
+                qml.RX(a[1], wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            return execute(
+                [tape],
+                dev,
+                gradient_fn="device",
+                mode="backward",
+                gradient_kwargs={"method": "adjoint_jacobian"},
+            )[0]
+
+        a = np.array([0.1, 0.2], requires_grad=True)
+        cost(a)
+
+        assert dev.num_executions == 1
+        spy_execute.assert_called()
+        spy_gradients.assert_not_called()
+
+        qml.jacobian(cost)(a)
+        spy_gradients.assert_called()
+
+
+execute_kwargs = [
+    {"gradient_fn": param_shift},
+    {
+        "gradient_fn": "device",
+        "mode": "forward",
+        "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True},
+    },
+    {
+        "gradient_fn": "device",
+        "mode": "backward",
+        "gradient_kwargs": {"method": "adjoint_jacobian"},
+    },
+]
+
+
+@pytest.mark.parametrize("execute_kwargs", execute_kwargs)
+class TestAutogradExecuteIntegration:
+    """Test the autograd interface execute function
+    integrates well for both forward and backward execution"""
+
+    def test_execution(self, execute_kwargs):
+        """Test execution"""
+        dev = qml.device("default.qubit", wires=1)
+
+        def cost(a, b):
+            with qml.tape.JacobianTape() as tape1:
+                qml.RY(a, wires=0)
+                qml.RX(b, wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            with qml.tape.JacobianTape() as tape2:
+                qml.RY(a, wires=0)
+                qml.RX(b, wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            return execute([tape1, tape2], dev, **execute_kwargs)
+
+        a = np.array(0.1, requires_grad=True)
+        b = np.array(0.2, requires_grad=False)
+        res = cost(a, b)
+
+        assert len(res) == 2
+        assert res[0].shape == (1,)
+        assert res[1].shape == (1,)
+
+    def test_scalar_jacobian(self, execute_kwargs, tol):
+        """Test scalar jacobian calculation"""
+        a = np.array(0.1, requires_grad=True)
+        dev = qml.device("default.qubit", wires=2)
+
+        def cost(a):
+            with qml.tape.JacobianTape() as tape:
+                qml.RY(a, wires=0)
+                qml.expval(qml.PauliZ(0))
+            return execute([tape], dev, **execute_kwargs)[0]
+
+        res = qml.jacobian(cost)(a)
+        assert res.shape == (1,)
+
+        # compare to standard tape jacobian
+        with qml.tape.JacobianTape() as tape:
+            qml.RY(a, wires=0)
+            qml.expval(qml.PauliZ(0))
+
+        tape.trainable_params = {0}
+        tapes, fn = param_shift(tape)
+        expected = fn(dev.batch_execute(tapes))
+
+        assert expected.shape == (1, 1)
+        assert np.allclose(res, np.squeeze(expected), atol=tol, rtol=0)
+
+    def test_jacobian(self, execute_kwargs, tol):
+        """Test jacobian calculation"""
+        a = np.array(0.1, requires_grad=True)
+        b = np.array(0.2, requires_grad=True)
+
+        def cost(a, b, device):
+            with qml.tape.JacobianTape() as tape:
+                qml.RY(a, wires=0)
+                qml.RX(b, wires=1)
+                qml.CNOT(wires=[0, 1])
+                qml.expval(qml.PauliZ(0))
+                qml.expval(qml.PauliY(1))
+            return execute([tape], device, **execute_kwargs)[0]
+
+        dev = qml.device("default.qubit", wires=2)
+
+        res = cost(a, b, device=dev)
+        expected = [np.cos(a), -np.cos(a) * np.sin(b)]
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        res = qml.jacobian(cost)(a, b, device=dev)
+        assert res.shape == (2, 2)
+
+        expected = [[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_reusing_quantum_tape(self, execute_kwargs, tol):
+        """Test re-using a quantum tape by passing new parameters"""
+        a = np.array(0.1, requires_grad=True)
+        b = np.array(0.2, requires_grad=True)
+
+        dev = qml.device("default.qubit", wires=2)
+
+        with qml.tape.JacobianTape() as tape:
+            qml.RY(a, wires=0)
+            qml.RX(b, wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0))
+            qml.expval(qml.PauliY(1))
+
+        assert tape.trainable_params == {0, 1}
+
+        def cost(a, b):
+            tape.set_parameters([a, b])
+            return execute([tape], dev, **execute_kwargs)[0]
+
+        jac_fn = qml.jacobian(cost)
+        jac = jac_fn(a, b)
+
+        a = np.array(0.54, requires_grad=True)
+        b = np.array(0.8, requires_grad=True)
+
+        # check that the cost function continues to depend on the
+        # values of the parameters for subsequent calls
+        res2 = cost(2 * a, b)
+        expected = [np.cos(2 * a), -np.cos(2 * a) * np.sin(b)]
+        assert np.allclose(res2, expected, atol=tol, rtol=0)
+
+        jac_fn = qml.jacobian(lambda a, b: cost(2 * a, b))
+        jac = jac_fn(a, b)
+        expected = [
+            [-2 * np.sin(2 * a), 0],
+            [2 * np.sin(2 * a) * np.sin(b), -np.cos(2 * a) * np.cos(b)],
+        ]
+        assert np.allclose(jac, expected, atol=tol, rtol=0)
+
+    def test_classical_processing(self, execute_kwargs, tol):
+        """Test classical processing within the quantum tape"""
+        a = np.array(0.1, requires_grad=True)
+        b = np.array(0.2, requires_grad=False)
+        c = np.array(0.3, requires_grad=True)
+
+        def cost(a, b, c, device):
+            with qml.tape.JacobianTape() as tape:
+                qml.RY(a * c, wires=0)
+                qml.RZ(b, wires=0)
+                qml.RX(c + c ** 2 + np.sin(a), wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            return execute([tape], device, **execute_kwargs)[0]
+
+        dev = qml.device("default.qubit", wires=2)
+        res = qml.jacobian(cost)(a, b, c, device=dev)
+        assert res.shape == (1, 2)
+
+    def test_no_trainable_parameters(self, execute_kwargs, tol):
+        """Test evaluation and Jacobian if there are no trainable parameters"""
+        a = np.array(0.1, requires_grad=False)
+        b = np.array(0.2, requires_grad=False)
+
+        def cost(a, b, device):
+            with qml.tape.JacobianTape() as tape:
+                qml.RY(a, wires=0)
+                qml.RX(b, wires=0)
+                qml.CNOT(wires=[0, 1])
+                qml.expval(qml.PauliZ(0))
+                qml.expval(qml.PauliZ(1))
+            return execute([tape], device, **execute_kwargs)[0]
+
+        dev = qml.device("default.qubit", wires=2)
+        res = cost(a, b, device=dev)
+        assert res.shape == (2,)
+
+        res = qml.jacobian(cost)(a, b, device=dev)
+        assert len(res) == 0
+
+        def loss(a, b):
+            return np.sum(cost(a, b, device=dev))
+
+        with pytest.warns(UserWarning, match="Output seems independent"):
+            res = qml.grad(loss)(a, b)
+
+        assert np.allclose(res, 0)
+
+    def test_matrix_parameter(self, execute_kwargs, tol):
+        """Test that the autograd interface works correctly
+        with a matrix parameter"""
+        U = np.array([[0, 1], [1, 0]], requires_grad=False)
+        a = np.array(0.1, requires_grad=True)
+
+        def cost(a, U, device):
+            with qml.tape.JacobianTape() as tape:
+                qml.QubitUnitary(U, wires=0)
+                qml.RY(a, wires=0)
+                qml.expval(qml.PauliZ(0))
+            return execute([tape], device, **execute_kwargs)[0]
+
+        dev = qml.device("default.qubit", wires=2)
+        res = cost(a, U, device=dev)
+        assert np.allclose(res, -np.cos(a), atol=tol, rtol=0)
+
+        jac_fn = qml.jacobian(cost)
+        res = jac_fn(a, U, device=dev)
+        assert np.allclose(res, np.sin(a), atol=tol, rtol=0)
+
+    def test_differentiable_expand(self, execute_kwargs, tol):
+        """Test that operation and nested tapes expansion
+        is differentiable"""
+
+        class U3(qml.U3):
+            def expand(self):
+                tape = qml.tape.JacobianTape()
+                theta, phi, lam = self.data
+                wires = self.wires
+                tape._ops += [
+                    qml.Rot(lam, theta, -lam, wires=wires),
+                    qml.PhaseShift(phi + lam, wires=wires),
+                ]
+                return tape
+
+        def cost_fn(a, p, device):
+            tape = qml.tape.JacobianTape()
+
+            with tape:
+                qml.RX(a, wires=0)
+                U3(*p, wires=0)
+                qml.expval(qml.PauliX(0))
+
+            tape = tape.expand(stop_at=lambda obj: device.supports_operation(obj.name))
+            return execute([tape], device, **execute_kwargs)[0]
+
+        a = np.array(0.1, requires_grad=False)
+        p = np.array([0.1, 0.2, 0.3], requires_grad=True)
+
+        dev = qml.device("default.qubit", wires=1)
+        res = cost_fn(a, p, device=dev)
+        expected = np.cos(a) * np.cos(p[1]) * np.sin(p[0]) + np.sin(a) * (
+            np.cos(p[2]) * np.sin(p[1]) + np.cos(p[0]) * np.cos(p[1]) * np.sin(p[2])
+        )
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        jac_fn = qml.jacobian(cost_fn)
+        res = jac_fn(a, p, device=dev)
+        expected = np.array(
+            [
+                np.cos(p[1]) * (np.cos(a) * np.cos(p[0]) - np.sin(a) * np.sin(p[0]) * np.sin(p[2])),
+                np.cos(p[1]) * np.cos(p[2]) * np.sin(a)
+                - np.sin(p[1])
+                * (np.cos(a) * np.sin(p[0]) + np.cos(p[0]) * np.sin(a) * np.sin(p[2])),
+                np.sin(a)
+                * (np.cos(p[0]) * np.cos(p[1]) * np.cos(p[2]) - np.sin(p[1]) * np.sin(p[2])),
+            ]
+        )
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_probability_differentiation(self, execute_kwargs, tol):
+        """Tests correct output shape and evaluation for a tape
+        with prob outputs"""
+
+        if execute_kwargs["gradient_fn"] == "device":
+            pytest.skip("Adjoint differentiation does not yet support probabilities")
+
+        def cost(x, y, device):
+            with qml.tape.JacobianTape() as tape:
+                qml.RX(x, wires=[0])
+                qml.RY(y, wires=[1])
+                qml.CNOT(wires=[0, 1])
+                qml.probs(wires=[0])
+                qml.probs(wires=[1])
+
+            return execute([tape], device, **execute_kwargs)[0]
+
+        dev = qml.device("default.qubit", wires=2)
+        x = np.array(0.543, requires_grad=True)
+        y = np.array(-0.654, requires_grad=True)
+
+        res = cost(x, y, device=dev)
+        expected = np.array(
+            [
+                [np.cos(x / 2) ** 2, np.sin(x / 2) ** 2],
+                [(1 + np.cos(x) * np.cos(y)) / 2, (1 - np.cos(x) * np.cos(y)) / 2],
+            ]
+        )
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        jac_fn = qml.jacobian(cost)
+        res = jac_fn(x, y, device=dev)
+        assert res.shape == (2, 2, 2)
+
+        expected = np.array(
+            [
+                [[-np.sin(x) / 2, 0], [np.sin(x) / 2, 0]],
+                [
+                    [-np.sin(x) * np.cos(y) / 2, -np.cos(x) * np.sin(y) / 2],
+                    [np.cos(y) * np.sin(x) / 2, np.cos(x) * np.sin(y) / 2],
+                ],
+            ]
+        )
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_ragged_differentiation(self, execute_kwargs, tol):
+        """Tests correct output shape and evaluation for a tape
+        with prob and expval outputs"""
+        if execute_kwargs["gradient_fn"] == "device":
+            pytest.skip("Adjoint differentiation does not yet support probabilities")
+
+        def cost(x, y, device):
+            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=[1])
+
+            return execute([tape], device, **execute_kwargs)[0]
+
+        dev = qml.device("default.qubit", wires=2)
+        x = np.array(0.543, requires_grad=True)
+        y = np.array(-0.654, requires_grad=True)
+
+        res = cost(x, y, device=dev)
+        expected = np.array(
+            [np.cos(x), (1 + np.cos(x) * np.cos(y)) / 2, (1 - np.cos(x) * np.cos(y)) / 2]
+        )
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        jac_fn = qml.jacobian(cost)
+        res = jac_fn(x, y, device=dev)
+        expected = np.array(
+            [
+                [-np.sin(x), 0],
+                [-np.sin(x) * np.cos(y) / 2, -np.cos(x) * np.sin(y) / 2],
+                [np.cos(y) * np.sin(x) / 2, np.cos(x) * np.sin(y) / 2],
+            ]
+        )
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_sampling(self, execute_kwargs):
+        """Test sampling works as expected"""
+        if execute_kwargs["gradient_fn"] == "device" and execute_kwargs["mode"] == "forward":
+            pytest.skip("Adjoint differentiation does not support samples")
+
+        def cost(x, device):
+            with qml.tape.JacobianTape() as tape:
+                qml.Hadamard(wires=[0])
+                qml.CNOT(wires=[0, 1])
+                qml.sample(qml.PauliZ(0))
+                qml.sample(qml.PauliX(1))
+
+            return execute([tape], device, **execute_kwargs)[0]
+
+        dev = qml.device("default.qubit", wires=2, shots=10)
+        x = np.array(0.543, requires_grad=True)
+        res = cost(x, device=dev)
+        assert res.shape == (2, 10)
+
+
+class TestHigherOrderDerivatives:
+    """Test that the autograd execute function can be differentiated"""
+
+    @pytest.mark.parametrize(
+        "params",
+        [
+            np.array([0.543, -0.654], requires_grad=True),
+            np.array([0, -0.654], requires_grad=True),
+            np.array([-2.0, 0], requires_grad=True),
+        ],
+    )
+    def test_parameter_shift_hessian(self, params, tol):
+        """Tests that the output of the parameter-shift 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 tape1:
+                qml.RX(x[0], wires=[0])
+                qml.RY(x[1], wires=[1])
+                qml.CNOT(wires=[0, 1])
+                qml.var(qml.PauliZ(0) @ qml.PauliX(1))
+
+            with qml.tape.JacobianTape() as tape2:
+                qml.RX(x[0], wires=0)
+                qml.RY(x[0], wires=1)
+                qml.CNOT(wires=[0, 1])
+                qml.probs(wires=1)
+
+            result = execute([tape1, tape2], dev, gradient_fn=param_shift)
+            return result[0] + result[1][0, 0]
+
+        res = cost_fn(params)
+        x, y = params
+        expected = 0.5 * (3 + np.cos(x) ** 2 * np.cos(2 * y))
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        res = qml.grad(cost_fn)(params)
+        expected = np.array(
+            [-np.cos(x) * np.cos(2 * y) * np.sin(x), -np.cos(x) ** 2 * np.sin(2 * y)]
+        )
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        res = qml.jacobian(qml.grad(cost_fn))(params)
+        expected = np.array(
+            [
+                [-np.cos(2 * x) * np.cos(2 * y), np.sin(2 * x) * np.sin(2 * y)],
+                [np.sin(2 * x) * np.sin(2 * y), -2 * np.cos(x) ** 2 * np.cos(2 * y)],
+            ]
+        )
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_adjoint_hessian(self, tol):
+        """Since the adjoint hessian is not a differentiable transform,
+        higher-order derivatives are not supported."""
+        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))
+
+            return execute(
+                [tape],
+                dev,
+                gradient_fn="device",
+                gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True},
+            )[0]
+
+        with pytest.warns(UserWarning, match="Output seems independent"):
+            res = qml.jacobian(qml.grad(cost_fn))(params)
+
+        assert np.allclose(res, np.zeros([2, 2]), atol=tol, rtol=0)