Repository: PennyLaneAI/pennylane
Pull request #3903 (2023-03-14 19:19:46)

Integrate new device interface with `qml.execute` and `autograd`
Now that `devices.experimental.DefaultQubit2` is added to master, we need to start integrating it with the workflow.

This PR adds support to `qml.execute` for the autograd interface for backprop and gradient transforms. The new device will only be supported with the new return types specification.

codecov[bot]: ## [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3903?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3903](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3903?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (ac53c5c) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/cd2a83867e5d4b97e70f3dbfbba0009d24ea8f27?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (cd2a838) will **decrease** coverage by `0.01%`.
> The diff coverage is `93.75%`.

```diff
@@            Coverage Diff             @@
##           master    #3903      +/-   ##
==========================================
- Coverage   99.72%   99.72%   -0.01%     
==========================================
  Files         351      351              
  Lines       30847    30872      +25     
==========================================
+ Hits        30763    30786      +23     
- Misses         84       86       +2     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3903?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/devices/experimental/execution\_config.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3903?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZXhwZXJpbWVudGFsL2V4ZWN1dGlvbl9jb25maWcucHk=) | `100.00% <ø> (ø)` | |
| [pennylane/interfaces/execution.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3903?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvZXhlY3V0aW9uLnB5) | `99.12% <92.85%> (-0.88%)` | :arrow_down: |
| [pennylane/interfaces/autograd.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3903?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYXV0b2dyYWQucHk=) | `100.00% <100.00%> (ø)` | |

... and [1 file with indirect coverage changes](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3903/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 6da4853f1..4dd765b11 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -78,6 +78,9 @@
   `use_device_gradient`, and `grad_on_execution`.
   [(#4102)](https://github.com/PennyLaneAI/pennylane/pull/4102)
 
+* The new device interface in integrated with `qml.execute` for autograd, backpropagation, and no differentiation.
+  [(#3903)](https://github.com/PennyLaneAI/pennylane/pull/3903)
+
 <h3>Breaking changes 💔</h3>
 
 * `DiagonalQubitUnitary` does not decompose into `QubitUnitary` any longer, but into `RZ`, `IsingZZ`
diff --git a/pennylane/devices/experimental/execution_config.py b/pennylane/devices/experimental/execution_config.py
index 248bf8e75..4a084872e 100644
--- a/pennylane/devices/experimental/execution_config.py
+++ b/pennylane/devices/experimental/execution_config.py
@@ -75,7 +75,7 @@ class ExecutionConfig:
         """
         if self.interface not in SUPPORTED_INTERFACES:
             raise ValueError(
-                f"interface must be in {SUPPORTED_INTERFACES}, got {self.interface} instead."
+                f"Unknown interface. interface must be in {SUPPORTED_INTERFACES}, got {self.interface} instead."
             )
 
         if self.grad_on_execution not in {True, False, None}:
diff --git a/pennylane/interfaces/autograd.py b/pennylane/interfaces/autograd.py
index 1b6c1651a..1c9e88c87 100644
--- a/pennylane/interfaces/autograd.py
+++ b/pennylane/interfaces/autograd.py
@@ -260,7 +260,6 @@ def _vjp_legacy(
         return_vjps = [
             qml.math.to_numpy(v, max_depth=_n) if isinstance(v, ArrayBox) else v for v in vjps
         ]
-
         return return_vjps
 
     return grad_fn
@@ -415,7 +414,12 @@ def vjp(
 
         jacs = []
         for t in tapes:
-            g_tapes, fn = gradient_fn(t, shots=device.shot_vector, **gradient_kwargs)
+            if isinstance(device, qml.devices.experimental.Device):  # pragma:  no-cover
+                # cant test until we integrate device with shot vector
+                shot_vector = t.shots.shot_vector if t.shots.has_partitioned_shots else None
+            else:
+                shot_vector = device.shot_vector
+            g_tapes, fn = gradient_fn(t, shots=shot_vector, **gradient_kwargs)
 
             unwrapped_tapes = tuple(convert_to_numpy_parameters(g_t) for g_t in g_tapes)
             res, _ = execute_fn(unwrapped_tapes, **gradient_kwargs)
@@ -433,6 +437,14 @@ def vjp(
 
         computing_jacobian = _n == max_diff
 
+        if isinstance(device, qml.devices.experimental.Device):  # pragma: no-cover
+            # assumes all tapes have the same shot vector
+            shot_vector = (
+                tapes[0].shots.shot_vector if tapes[0].shots.has_partitioned_shots else None
+            )
+        else:
+            shot_vector = device.shot_vector
+
         if gradient_fn and gradient_fn.__name__ == "param_shift" and computing_jacobian:
             jacs = _get_jac_with_caching()
         else:
@@ -441,7 +453,7 @@ def vjp(
         if jacs:
             # Jacobians were computed on the forward pass (mode="forward") or the Jacobian was cached
             # No additional quantum evaluations needed; simply compute the VJPs directly.
-            vjps = _compute_vjps_autograd(jacs, dy, multi_measurements, device.shot_vector)
+            vjps = _compute_vjps_autograd(jacs, dy, multi_measurements, shot_vector)
 
         else:
             # Need to compute the Jacobians on the backward pass (accumulation="backward")
@@ -455,7 +467,7 @@ def vjp(
                         unwrapped_tapes,
                         dy,
                         gradient_fn,
-                        shots=device.shot_vector,
+                        shots=shot_vector,
                         reduction="append",
                         gradient_kwargs=gradient_kwargs,
                     )
@@ -467,7 +479,7 @@ def vjp(
                         tapes,
                         dy,
                         gradient_fn,
-                        shots=device.shot_vector,
+                        shots=shot_vector,
                         reduction="append",
                         gradient_kwargs=gradient_kwargs,
                     )
diff --git a/pennylane/interfaces/execution.py b/pennylane/interfaces/execution.py
index 7a0602a28..fde3477cf 100644
--- a/pennylane/interfaces/execution.py
+++ b/pennylane/interfaces/execution.py
@@ -27,15 +27,17 @@ import inspect
 import warnings
 from contextlib import _GeneratorContextManager
 from functools import wraps, partial
-from typing import Callable, Sequence
+from typing import Callable, Sequence, Optional, Union, Tuple
 
-from cachetools import LRUCache
+from cachetools import LRUCache, Cache
 
 import pennylane as qml
 from pennylane.tape import QuantumTape
 
 from .set_shots import set_shots
 
+device_type = Union[qml.Device, "qml.devices.experimental.Device"]
+
 INTERFACE_MAP = {
     None: "Numpy",
     "auto": "auto",
@@ -85,6 +87,83 @@ def _adjoint_jacobian_expansion(
     return tapes
 
 
+def _batch_transform(
+    tapes: Sequence[QuantumTape],
+    device: device_type,
+    config: "qml.devices.experimental.ExecutionConfig",
+    override_shots: Union[bool, int, Sequence[int]] = False,
+    device_batch_transform: bool = True,
+) -> Tuple[Sequence[QuantumTape], Callable, "qml.devices.experimental.ExecutionConfig"]:
+    """Apply the device batch transform unless requested not to.
+
+    Args:
+        tapes (Tuple[.QuantumTape]): batch of tapes to preprocess
+        device (Device, devices.experimental.Device): the device that defines the required batch transformation
+        config (qml.devices.experimental.ExecutionConfig): the config that characterizes the requested computation
+        override_shots (int): The number of shots to use for the execution. If ``False``, then the
+            number of shots on the device is used.
+        device_batch_transform (bool): Whether to apply any batch transforms defined by the device
+            (within :meth:`Device.batch_transform`) to each tape to be executed. The default behaviour
+            of the device batch transform is to expand out Hamiltonian measurements into
+            constituent terms if not supported on the device.
+
+    Returns:
+        Sequence[QuantumTape], Callable: The new batch of quantum scripts and the post processing
+
+    """
+    if isinstance(device, qml.devices.experimental.Device):
+        if not device_batch_transform:
+            warnings.warn(
+                "device batch transforms cannot be turned off with the new device interface.",
+                UserWarning,
+            )
+        return device.preprocess(tapes, config)
+    if device_batch_transform:
+        dev_batch_transform = set_shots(device, override_shots)(device.batch_transform)
+        return *qml.transforms.map_batch_transform(dev_batch_transform, tapes), config
+
+    def null_post_processing_fn(results):
+        """A null post processing function used because the user requested not to use the device batch transform."""
+        return results
+
+    return tapes, null_post_processing_fn, config
+
+
+def _preprocess_expand_fn(
+    expand_fn: Union[str, Callable], device: device_type, max_expansion: int
+) -> Callable:
+    """Preprocess the ``expand_fn`` configuration property.
+
+    Args:
+        expand_fn (str, Callable): If string, then it must be "device".  Otherwise, it should be a map
+            from one tape to a new tape. The final tape must be natively executable by the device.
+        device (Device, devices.experimental.Device): The device that we will be executing on.
+        max_expansion (int): The number of times the internal circuit should be expanded when
+            executed on a device. Expansion occurs when an operation or measurement is not
+            supported, and results in a gate decomposition. If any operations in the decomposition
+            remain unsupported by the device, another expansion occurs.
+
+    Returns:
+        Callable: a map from one quantum tape to a new one. The output should be compatible with the device.
+
+    """
+    if expand_fn != "device":
+        return expand_fn
+    if isinstance(device, qml.devices.experimental.Device):
+
+        def blank_expansion_function(tape):  # pylint: disable=function-redefined
+            """A blank expansion function since the new device handles expansion in preprocessing."""
+            return tape
+
+        return blank_expansion_function
+
+    def device_expansion_function(tape):  # pylint: disable=function-redefined
+        """A wrapper around the device ``expand_fn``."""
+        return device.expand_fn(tape, max_expansion=max_expansion)
+
+    return device_expansion_function
+
+
 def cache_execute(fn: Callable, cache, pass_kwargs=False, return_tuple=True, expand_fn=None):
     """Decorator that adds caching to a function that executes
     multiple tapes on a device.
@@ -137,7 +216,9 @@ def cache_execute(fn: Callable, cache, pass_kwargs=False, return_tuple=True, exp
         if cache is None or (isinstance(cache, bool) and not cache):
             # No caching. Simply execute the execution function
             # and return the results.
-            res = fn(tapes, **kwargs)
+
+            # must convert to list as new device interface returns tuples
+            res = list(fn(tapes, **kwargs))
             return (res, []) if return_tuple else res
 
         execution_tapes = {}
@@ -201,7 +282,8 @@ def cache_execute(fn: Callable, cache, pass_kwargs=False, return_tuple=True, exp
 
         else:
             # execute all unique tapes that do not exist in the cache
-            res = fn(execution_tapes.values(), **kwargs)
+            # convert to list as new device interface returns a tuple
+            res = list(fn(execution_tapes.values(), **kwargs))
 
         final_res = []
 
@@ -228,16 +310,16 @@ def cache_execute(fn: Callable, cache, pass_kwargs=False, return_tuple=True, exp
 
 def execute(
     tapes: Sequence[QuantumTape],
-    device,
-    gradient_fn: Callable = None,
+    device: device_type,
+    gradient_fn: Optional[Union[Callable, str]] = None,
     interface="auto",
     grad_on_execution="best",
     gradient_kwargs=None,
-    cache=True,
+    cache: Union[bool, dict, Cache] = True,
     cachesize=10000,
     max_diff=1,
     override_shots: int = False,
-    expand_fn="device",
+    expand_fn="device",  # type: ignore
     max_expansion=10,
     device_batch_transform=True,
 ):
@@ -261,7 +343,7 @@ def execute(
             pass. The 'best' option chooses automatically between the two options and is default.
         gradient_kwargs (dict): dictionary of keyword arguments to pass when
             determining the gradients of tapes
-        cache (bool): Whether to cache evaluations. This can result in
+        cache (bool, dict, Cache): Whether to cache evaluations. This can result in
             a significant reduction in quantum evaluations during gradient computations.
         cachesize (int): the size of the cache
         max_diff (int): If ``gradient_fn`` is a gradient transform, this option specifies
@@ -270,7 +352,7 @@ def execute(
             (classical) computational overhead during the backwards pass.
         override_shots (int): The number of shots to use for the execution. If ``False``, then the
             number of shots on the device is used.
-        expand_fn (function): Tape expansion function to be called prior to device execution.
+        expand_fn (str, function): Tape expansion function to be called prior to device execution.
             Must have signature of the form ``expand_fn(tape, max_expansion)``, and return a
             single :class:`~.QuantumTape`. If not provided, by default :meth:`Device.expand_fn`
             is called.
@@ -363,53 +445,62 @@ def execute(
             device_batch_transform=device_batch_transform,
         )
 
+    ### Specifying and preprocessing variables ####
+
     if interface == "auto":
         params = []
         for tape in tapes:
             params.extend(tape.get_parameters(trainable_only=False))
         interface = qml.math.get_interface(*params)
 
+    new_device_interface = isinstance(device, qml.devices.experimental.Device)
+    config = qml.devices.experimental.ExecutionConfig(interface=interface)
     gradient_kwargs = gradient_kwargs or {}
 
-    if device_batch_transform:
-        dev_batch_transform = set_shots(device, override_shots)(device.batch_transform)
-        tapes, batch_fn = qml.transforms.map_batch_transform(dev_batch_transform, tapes)
-    else:
-        batch_fn = lambda res: res  # pragma: no cover
-
     if isinstance(cache, bool) and cache:
         # cache=True: create a LRUCache object
         cache = LRUCache(maxsize=cachesize)
         setattr(cache, "_persistent_cache", False)
 
-    batch_execute = set_shots(device, override_shots)(device.batch_execute)
+    if new_device_interface:
+        batch_execute = device.execute
+    else:
+        batch_execute = set_shots(device, override_shots)(device.batch_execute)
 
-    if expand_fn == "device":
-        expand_fn = lambda tape: device.expand_fn(tape, max_expansion=max_expansion)
+    expand_fn = _preprocess_expand_fn(expand_fn, device, max_expansion)
 
-    if gradient_fn is None:
-        # don't unwrap if it's an interface device
-        if "passthru_interface" in device.capabilities() or device.short_name == "default.mixed":
-            return batch_fn(
-                qml.interfaces.cache_execute(
-                    batch_execute, cache, return_tuple=False, expand_fn=expand_fn
-                )(tapes)
-            )
-        unwrapped_tapes = tuple(qml.transforms.convert_to_numpy_parameters(t) for t in tapes)
-        res = qml.interfaces.cache_execute(
-            batch_execute, cache, return_tuple=False, expand_fn=expand_fn
-        )(unwrapped_tapes)
+    #### Executing the configured setup #####
 
-        return batch_fn(res)
+    tapes, batch_fn, config = _batch_transform(
+        tapes, device, config, override_shots, device_batch_transform
+    )
 
-    if gradient_fn == "backprop" or interface is None:
-        return batch_fn(
-            qml.interfaces.cache_execute(
-                batch_execute, cache, return_tuple=False, expand_fn=expand_fn
-            )(tapes)
+    # Exiting early if we do not need to deal with an interface boundary
+    no_interface_boundary_required = interface is None or gradient_fn in {None, "backprop"}
+    if no_interface_boundary_required:
+        device_supports_interface_data = (
+            new_device_interface
+            or config.interface is None
+            or gradient_fn == "backprop"
+            or device.short_name == "default.mixed"
+            or "passthru_interface" in device.capabilities()
+        )
+        if not device_supports_interface_data:
+            tapes = tuple(qml.transforms.convert_to_numpy_parameters(t) for t in tapes)
+
+        # use qml.interfaces so that mocker can spy on it during testing
+        cached_execute_fn = qml.interfaces.cache_execute(
+            batch_execute,
+            cache,
+            expand_fn=expand_fn,
+            return_tuple=False,
+            pass_kwargs=new_device_interface,
         )
+        results = cached_execute_fn(tapes, execution_config=config)
+        return batch_fn(results)
 
     # the default execution function is batch_execute
+    # use qml.interfaces so that mocker can spy on it during testing
     execute_fn = qml.interfaces.cache_execute(batch_execute, cache, expand_fn=expand_fn)
 
     _grad_on_execution = False
@@ -438,11 +529,14 @@ def execute(
 
         else:
             # disable caching on the forward pass
+            # use qml.interfaces so that mocker can spy on it during testing
             execute_fn = qml.interfaces.cache_execute(batch_execute, cache=None)
 
             # replace the backward gradient computation
+            # use qml.interfaces so that mocker can spy on it during testing
+            gradient_fn_with_shots = set_shots(device, override_shots)(device.gradients)
             gradient_fn = qml.interfaces.cache_execute(
-                set_shots(device, override_shots)(device.gradients),
+                gradient_fn_with_shots,
                 cache,
                 pass_kwargs=True,
                 return_tuple=False,
@@ -453,12 +547,7 @@ def execute(
         # in this case would have ambiguous behaviour.
         raise ValueError("Gradient transforms cannot be used with grad_on_execution=True")
 
-    try:
-        mapped_interface = INTERFACE_MAP[interface]
-    except KeyError as e:
-        raise ValueError(
-            f"Unknown interface {interface}. Supported " f"interfaces are {SUPPORTED_INTERFACES}"
-        ) from e
+    mapped_interface = INTERFACE_MAP[config.interface]
     try:
         if mapped_interface == "autograd":
             from .autograd import execute as _execute
@@ -495,7 +584,7 @@ def execute(
 
 def _execute_legacy(
     tapes: Sequence[QuantumTape],
-    device,
+    device: device_type,
     gradient_fn: Callable = None,
     interface="auto",
     mode="best",
@@ -608,6 +697,10 @@ def _execute_legacy(
            [ 0.01983384, -0.97517033,  0.        ],
            [ 0.        ,  0.        , -0.95533649]])
     """
+
+    if isinstance(device, qml.devices.experimental.Device):
+        raise ValueError("New device interface only works with return types enabled.")
+
     if interface == "auto":
         params = []
         for tape in tapes:


Test patch:
diff --git a/tests/devices/qubit/test_preprocess.py b/tests/devices/qubit/test_preprocess.py
index 9bad1a2b3..cacfbf4d5 100644
--- a/tests/devices/qubit/test_preprocess.py
+++ b/tests/devices/qubit/test_preprocess.py
@@ -149,20 +149,15 @@ class TestExpandFnValidation:
         """Test that expand_fn throws an error when a tensor includes invalid obserables"""
         tape = QuantumScript(
             ops=[qml.PauliX(0), qml.PauliY(1)],
-            measurements=[
-                qml.expval(qml.GellMann(wires=0, index=1) @ qml.GellMann(wires=1, index=2))
-            ],
+            measurements=[qml.expval(qml.PauliX(0) @ qml.GellMann(wires=1, index=2))],
         )
         with pytest.raises(DeviceError, match="Observable expval"):
             expand_fn(tape)
 
-    def test_expand_fn_valid_tensor_observable(self):
-        """Test that expand_fn passes when a tensor includes only valid obserables."""
-        tape = QuantumScript(
-            ops=[qml.PauliX(0), qml.PauliY(1)],
-            measurements=[qml.expval(qml.PauliZ(0) @ qml.PauliX(1))],
-        )
-        expand_fn(tape)
+    def test_valid_tensor_observable(self):
+        """Test that a valid tensor ovservable passes without error."""
+        tape = QuantumScript([], [qml.expval(qml.PauliZ(0) @ qml.PauliY(1))])
+        assert expand_fn(tape) is tape
 
     def test_state_prep_only_one(self):
         """Test that a device error is raised if the script has multiple state prep ops."""
diff --git a/tests/interfaces/default_qubit_2_integration/test_autograd_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_autograd_default_qubit_2.py
new file mode 100644
index 000000000..09645cf42
--- /dev/null
+++ b/tests/interfaces/default_qubit_2_integration/test_autograd_default_qubit_2.py
@@ -0,0 +1,689 @@
+# Copyright 2018-2023 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.
+"""Autograd specific tests for execute and default qubit 2."""
+# pylint: disable=invalid-sequence-index
+import autograd
+import pytest
+from pennylane import numpy as np
+
+import pennylane as qml
+from pennylane.devices.experimental import DefaultQubit2
+from pennylane.gradients import param_shift
+from pennylane.interfaces import execute
+
+pytestmark = pytest.mark.autograd
+
+
+# pylint: disable=too-few-public-methods
+class TestCaching:
+    """Tests for caching behaviour"""
+
+    @pytest.mark.parametrize("num_params", [2, 3])
+    def test_caching_param_shift_hessian(self, num_params):
+        """Test that, when using parameter-shift transform,
+        caching reduces the number of evaluations to their optimum
+        when computing Hessians."""
+        dev = DefaultQubit2()
+        params = np.arange(1, num_params + 1) / 10
+
+        N = len(params)
+
+        def cost(x, cache):
+            with qml.queuing.AnnotatedQueue() as q:
+                qml.RX(x[0], wires=[0])
+                qml.RY(x[1], wires=[1])
+
+                for i in range(2, num_params):
+                    qml.RZ(x[i], wires=[i % 2])
+
+                qml.CNOT(wires=[0, 1])
+                qml.var(qml.prod(qml.PauliZ(0), qml.PauliX(1)))
+
+            tape = qml.tape.QuantumScript.from_queue(q)
+            return qml.execute(
+                [tape], dev, gradient_fn=qml.gradients.param_shift, cache=cache, max_diff=2
+            )[0]
+
+        # No caching: number of executions is not ideal
+        with qml.Tracker(dev) as tracker:
+            hess1 = qml.jacobian(qml.grad(cost))(params, cache=False)
+
+        if num_params == 2:
+            # compare to theoretical result
+            x, y, *_ = params
+            expected = np.array(
+                [
+                    [2 * np.cos(2 * x) * np.sin(y) ** 2, 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(expected, hess1)
+
+        expected_runs = 1  # forward pass
+
+        # Jacobian of an involutory observable:
+        # ------------------------------------
+        #
+        # 2 * N execs: evaluate the analytic derivative of <A>
+        # 1 execs: Get <A>, the expectation value of the tape with unshifted parameters.
+        num_shifted_evals = 2 * N
+        runs_for_jacobian = num_shifted_evals + 1
+        expected_runs += runs_for_jacobian
+
+        # Each tape used to compute the Jacobian is then shifted again
+        expected_runs += runs_for_jacobian * num_shifted_evals
+        assert tracker.totals["executions"] == expected_runs
+
+        # Use caching: number of executions is ideal
+
+        with qml.Tracker(dev) as tracker2:
+            hess2 = qml.jacobian(qml.grad(cost))(params, cache=True)
+        assert np.allclose(hess1, hess2)
+
+        expected_runs_ideal = 1  # forward pass
+        expected_runs_ideal += 2 * N  # Jacobian
+        expected_runs_ideal += N + 1  # Hessian diagonal
+        expected_runs_ideal += 4 * N * (N - 1) // 2  # Hessian off-diagonal
+        assert tracker2.totals["executions"] == expected_runs_ideal
+        assert expected_runs_ideal < expected_runs
+
+
+execute_kwargs_iterable = [
+    {"gradient_fn": param_shift},
+    {"gradient_fn": "backprop"},
+    # no device gradient yet
+]
+
+
+@pytest.mark.parametrize("execute_kwargs", execute_kwargs_iterable)
+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 = DefaultQubit2()
+
+        def cost(a, b):
+            ops1 = [qml.RY(a, wires=0), qml.RX(b, wires=0)]
+            tape1 = qml.tape.QuantumScript(ops1, [qml.expval(qml.PauliZ(0))])
+
+            ops2 = [qml.RY(a, wires="a"), qml.RX(b, wires="a")]
+            tape2 = qml.tape.QuantumScript(ops2, [qml.expval(qml.PauliZ("a"))])
+
+            return execute([tape1, tape2], dev, **execute_kwargs)
+
+        a = np.array(0.1, requires_grad=True)
+        b = np.array(0.2, requires_grad=False)
+        with dev.tracker:
+            res = cost(a, b)
+
+        assert dev.tracker.totals["batches"] == 1
+        assert dev.tracker.totals["executions"] == 2  # different wires so different hashes
+
+        assert len(res) == 2
+        assert res[0].shape == ()
+        assert res[1].shape == ()
+
+        assert qml.math.allclose(res[0], np.cos(a) * np.cos(b))
+        assert qml.math.allclose(res[1], np.cos(a) * np.cos(b))
+
+    def test_scalar_jacobian(self, execute_kwargs, tol):
+        """Test scalar jacobian calculation"""
+        a = np.array(0.1, requires_grad=True)
+        dev = DefaultQubit2()
+
+        def cost(a):
+            tape = qml.tape.QuantumScript([qml.RY(a, 0)], [qml.expval(qml.PauliZ(0))])
+            return execute([tape], dev, **execute_kwargs)[0]
+
+        res = qml.jacobian(cost)(a)
+        assert res.shape == ()  # pylint: disable=no-member
+
+        # compare to standard tape jacobian
+        tape = qml.tape.QuantumScript([qml.RY(a, wires=0)], [qml.expval(qml.PauliZ(0))])
+        tape.trainable_params = [0]
+        tapes, fn = param_shift(tape)
+        expected = fn(dev.execute(tapes))
+
+        assert expected.shape == ()
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+        assert np.allclose(res, -np.sin(a))
+
+    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):
+            ops = [qml.RY(a, wires=0), qml.RX(b, wires=1), qml.CNOT(wires=[0, 1])]
+            m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))]
+            tape = qml.tape.QuantumScript(ops, m)
+            return autograd.numpy.hstack(execute([tape], device, **execute_kwargs)[0])
+
+        dev = DefaultQubit2()
+
+        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 isinstance(res, tuple) and len(res) == 2
+        assert res[0].shape == (2,)
+        assert res[1].shape == (2,)
+
+        expected = ([-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)])
+        assert all(np.allclose(_r, _e, atol=tol, rtol=0) for _r, _e in zip(res, expected))
+
+    @pytest.mark.filterwarnings("ignore:Attempted to compute the gradient")
+    def test_tape_no_parameters(self, execute_kwargs, tol):
+        """Test that a tape with no parameters is correctly
+        ignored during the gradient computation"""
+
+        if execute_kwargs["gradient_fn"] == "device":
+            pytest.skip("Adjoint differentiation does not yet support probabilities")
+
+        dev = DefaultQubit2()
+
+        def cost(params):
+            tape1 = qml.tape.QuantumScript([qml.Hadamard(0)], [qml.expval(qml.PauliX(0))])
+
+            tape2 = qml.tape.QuantumScript(
+                [qml.RY(np.array(0.5, requires_grad=False), wires=0)], [qml.expval(qml.PauliZ(0))]
+            )
+
+            tape3 = qml.tape.QuantumScript(
+                [qml.RY(params[0], 0), qml.RX(params[1], 0)], [qml.expval(qml.PauliZ(0))]
+            )
+
+            tape4 = qml.tape.QuantumScript(
+                [qml.RY(np.array(0.5, requires_grad=False), 0)], [qml.probs(wires=[0, 1])]
+            )
+            return sum(
+                autograd.numpy.hstack(execute([tape1, tape2, tape3, tape4], dev, **execute_kwargs))
+            )
+
+        params = np.array([0.1, 0.2], requires_grad=True)
+        x, y = params
+
+        res = cost(params)
+        expected = 2 + np.cos(0.5) + np.cos(x) * np.cos(y)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        grad = qml.grad(cost)(params)
+        expected = [-np.cos(y) * np.sin(x), -np.cos(x) * np.sin(y)]
+        assert np.allclose(grad, expected, atol=tol, rtol=0)
+
+    @pytest.mark.filterwarnings("ignore:Attempted to compute the gradient")
+    def test_tapes_with_different_return_size(self, execute_kwargs):
+        """Test that tapes wit different can be executed and differentiated."""
+
+        if execute_kwargs["gradient_fn"] == "backprop":
+            pytest.xfail("backprop is not compatible with something about this situation.")
+
+        dev = DefaultQubit2()
+
+        def cost(params):
+            tape1 = qml.tape.QuantumScript(
+                [qml.RY(params[0], 0), qml.RX(params[1], 0)],
+                [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))],
+            )
+
+            tape2 = qml.tape.QuantumScript(
+                [qml.RY(np.array(0.5, requires_grad=False), 0)], [qml.expval(qml.PauliZ(0))]
+            )
+
+            tape3 = qml.tape.QuantumScript(
+                [qml.RY(params[0], 0), qml.RX(params[1], 0)], [qml.expval(qml.PauliZ(0))]
+            )
+            return autograd.numpy.hstack(execute([tape1, tape2, tape3], dev, **execute_kwargs))
+
+        params = np.array([0.1, 0.2], requires_grad=True)
+        x, y = params
+
+        res = cost(params)
+        assert isinstance(res, np.ndarray)
+        assert res.shape == (4,)
+
+        assert np.allclose(res[0], np.cos(x) * np.cos(y))
+        assert np.allclose(res[1], 1)
+        assert np.allclose(res[2], np.cos(0.5))
+        assert np.allclose(res[3], np.cos(x) * np.cos(y))
+
+        jac = qml.jacobian(cost)(params)
+        assert isinstance(jac, np.ndarray)
+        assert jac.shape == (4, 2)  # pylint: disable=no-member
+
+        assert np.allclose(jac[1:3], 0)
+
+        d1 = -np.sin(x) * np.cos(y)
+        assert np.allclose(jac[0, 0], d1)
+        assert np.allclose(jac[3, 0], d1)
+
+        d2 = -np.cos(x) * np.sin(y)
+        assert np.allclose(jac[0, 1], d2)
+        assert np.allclose(jac[3, 1], d2)
+
+    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 = DefaultQubit2()
+
+        tape = qml.tape.QuantumScript(
+            [qml.RY(a, 0), qml.RX(b, 1), qml.CNOT((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 autograd.numpy.hstack(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), 2 * np.sin(2 * a) * np.sin(b)],
+            [0, -np.cos(2 * a) * np.cos(b)],
+        )
+        assert isinstance(jac, tuple) and len(jac) == 2
+        assert all(np.allclose(_j, _e, atol=tol, rtol=0) for _j, _e in zip(jac, expected))
+
+    def test_classical_processing(self, execute_kwargs):
+        """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):
+            ops = [
+                qml.RY(a * c, wires=0),
+                qml.RZ(b, wires=0),
+                qml.RX(c + c**2 + np.sin(a), wires=0),
+            ]
+
+            tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))])
+            return execute([tape], device, **execute_kwargs)[0]
+
+        dev = DefaultQubit2()
+        res = qml.jacobian(cost)(a, b, c, device=dev)
+
+        # Only two arguments are trainable
+        assert isinstance(res, tuple) and len(res) == 2
+        assert res[0].shape == ()
+        assert res[1].shape == ()
+
+        # I tried getting analytic results for this circuit but I kept being wrong and am giving up
+
+    def test_no_trainable_parameters(self, execute_kwargs):
+        """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):
+            ops = [qml.RY(a, 0), qml.RX(b, 0), qml.CNOT((0, 1))]
+            m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))]
+            tape = qml.tape.QuantumScript(ops, m)
+            return autograd.numpy.hstack(execute([tape], device, **execute_kwargs)[0])
+
+        dev = DefaultQubit2()
+        res = cost(a, b, device=dev)
+        assert res.shape == (2,)
+
+        with pytest.warns(UserWarning, match="Attempted to differentiate a function with no"):
+            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="Attempted to differentiate a function with no"):
+            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):
+            ops = [qml.QubitUnitary(U, wires=0), qml.RY(a, wires=0)]
+            tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))])
+            return execute([tape], device, **execute_kwargs)[0]
+
+        dev = DefaultQubit2()
+        res = cost(a, U, device=dev)
+        assert np.allclose(res, -np.cos(a), atol=tol, rtol=0)
+
+        jac_fn = qml.jacobian(cost)
+        jac = jac_fn(a, U, device=dev)
+        assert isinstance(jac, np.ndarray)
+        assert np.allclose(jac, 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):
+            """Dummy operator."""
+
+            def expand(self):
+                theta, phi, lam = self.data
+                wires = self.wires
+                return qml.tape.QuantumScript(
+                    [
+                        qml.Rot(lam, theta, -lam, wires=wires),
+                        qml.PhaseShift(phi + lam, wires=wires),
+                    ]
+                )
+
+        def cost_fn(a, p, device):
+            tape = qml.tape.QuantumScript(
+                [qml.RX(a, wires=0), U3(*p, wires=0)], [qml.expval(qml.PauliX(0))]
+            )
+            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 = DefaultQubit2()
+        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"""
+
+        def cost(x, y, device):
+            ops = [qml.RX(x, 0), qml.RY(y, 1), qml.CNOT((0, 1))]
+            m = [qml.probs(wires=0), qml.probs(wires=1)]
+            tape = qml.tape.QuantumScript(ops, m)
+            return autograd.numpy.hstack(execute([tape], device, **execute_kwargs)[0])
+
+        dev = DefaultQubit2()
+        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 isinstance(res, tuple) and len(res) == 2
+        assert res[0].shape == (4,)
+        assert res[1].shape == (4,)
+
+        expected = (
+            np.array(
+                [
+                    [
+                        -np.sin(x) / 2,
+                        np.sin(x) / 2,
+                        -np.sin(x) * np.cos(y) / 2,
+                        np.sin(x) * np.cos(y) / 2,
+                    ],
+                ]
+            ),
+            np.array(
+                [
+                    [0, 0, -np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2],
+                ]
+            ),
+        )
+
+        assert np.allclose(res[0], expected[0], atol=tol, rtol=0)
+        assert np.allclose(res[1], expected[1], 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):
+            ops = [qml.RX(x, wires=0), qml.RY(y, 1), qml.CNOT((0, 1))]
+            m = [qml.expval(qml.PauliZ(0)), qml.probs(wires=1)]
+            tape = qml.tape.QuantumScript(ops, m)
+            return autograd.numpy.hstack(execute([tape], device, **execute_kwargs)[0])
+
+        dev = DefaultQubit2()
+        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)
+        assert isinstance(res, tuple) and len(res) == 2
+        assert res[0].shape == (3,)
+        assert res[1].shape == (3,)
+
+        expected = (
+            np.array([-np.sin(x), -np.sin(x) * np.cos(y) / 2, np.sin(x) * np.cos(y) / 2]),
+            np.array([0, -np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2]),
+        )
+        assert np.allclose(res[0], expected[0], atol=tol, rtol=0)
+        assert np.allclose(res[1], expected[1], atol=tol, rtol=0)
+
+
+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 = DefaultQubit2()
+
+        def cost_fn(x):
+            ops1 = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))]
+            tape1 = qml.tape.QuantumScript(ops1, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))])
+
+            ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))]
+            tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)])
+            result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=2)
+            return result[0] + result[1][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_max_diff(self, tol):
+        """Test that setting the max_diff parameter blocks higher-order
+        derivatives"""
+        dev = DefaultQubit2()
+        params = np.array([0.543, -0.654], requires_grad=True)
+
+        def cost_fn(x):
+            ops = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))]
+            tape1 = qml.tape.QuantumScript(ops, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))])
+
+            ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))]
+            tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)])
+
+            result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=1)
+            return result[0] + result[1][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)
+
+        with pytest.warns(UserWarning, match="Output seems independent"):
+            res = qml.jacobian(qml.grad(cost_fn))(params)
+
+        expected = np.zeros([2, 2])
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+
+@pytest.mark.parametrize("execute_kwargs", execute_kwargs_iterable)
+class TestHamiltonianWorkflows:
+    """Test that tapes ending with expectations
+    of Hamiltonians provide correct results and gradients"""
+
+    @pytest.fixture
+    def cost_fn(self, execute_kwargs):
+        """Cost function for gradient tests"""
+
+        def _cost_fn(weights, coeffs1, coeffs2, dev=None):
+            obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)]
+            H1 = qml.Hamiltonian(coeffs1, obs1)
+
+            obs2 = [qml.PauliZ(0)]
+            H2 = qml.Hamiltonian(coeffs2, obs2)
+
+            with qml.queuing.AnnotatedQueue() as q:
+                qml.RX(weights[0], wires=0)
+                qml.RY(weights[1], wires=1)
+                qml.CNOT(wires=[0, 1])
+                qml.expval(H1)
+                qml.expval(H2)
+
+            tape = qml.tape.QuantumScript.from_queue(q)
+            return autograd.numpy.hstack(execute([tape], dev, **execute_kwargs)[0])
+
+        return _cost_fn
+
+    @staticmethod
+    def cost_fn_expected(weights, coeffs1, coeffs2):
+        """Analytic value of cost_fn above"""
+        a, b, c = coeffs1
+        d = coeffs2[0]
+        x, y = weights
+        return [-c * np.sin(x) * np.sin(y) + np.cos(x) * (a + b * np.sin(y)), d * np.cos(x)]
+
+    @staticmethod
+    def cost_fn_jacobian(weights, coeffs1, coeffs2):
+        """Analytic jacobian of cost_fn above"""
+        a, b, c = coeffs1
+        d = coeffs2[0]
+        x, y = weights
+        return np.array(
+            [
+                [
+                    -c * np.cos(x) * np.sin(y) - np.sin(x) * (a + b * np.sin(y)),
+                    b * np.cos(x) * np.cos(y) - c * np.cos(y) * np.sin(x),
+                    np.cos(x),
+                    np.cos(x) * np.sin(y),
+                    -(np.sin(x) * np.sin(y)),
+                    0,
+                ],
+                [-d * np.sin(x), 0, 0, 0, 0, np.cos(x)],
+            ]
+        )
+
+    def test_multiple_hamiltonians_not_trainable(self, cost_fn, tol):
+        """Test hamiltonian with no trainable parameters."""
+        coeffs1 = np.array([0.1, 0.2, 0.3], requires_grad=False)
+        coeffs2 = np.array([0.7], requires_grad=False)
+        weights = np.array([0.4, 0.5], requires_grad=True)
+        dev = DefaultQubit2()
+
+        res = cost_fn(weights, coeffs1, coeffs2, dev=dev)
+        expected = self.cost_fn_expected(weights, coeffs1, coeffs2)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        res = qml.jacobian(cost_fn)(weights, coeffs1, coeffs2, dev=dev)
+        expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)[:, :2]
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_multiple_hamiltonians_trainable(self, cost_fn, tol):
+        """Test hamiltonian with trainable parameters."""
+        coeffs1 = np.array([0.1, 0.2, 0.3], requires_grad=True)
+        coeffs2 = np.array([0.7], requires_grad=True)
+        weights = np.array([0.4, 0.5], requires_grad=True)
+        dev = DefaultQubit2()
+
+        res = cost_fn(weights, coeffs1, coeffs2, dev=dev)
+        expected = self.cost_fn_expected(weights, coeffs1, coeffs2)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        res = np.hstack(qml.jacobian(cost_fn)(weights, coeffs1, coeffs2, dev=dev))
+        expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
diff --git a/tests/interfaces/default_qubit_2_integration/test_execute_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_execute_default_qubit_2.py
new file mode 100644
index 000000000..85a6b71fc
--- /dev/null
+++ b/tests/interfaces/default_qubit_2_integration/test_execute_default_qubit_2.py
@@ -0,0 +1,189 @@
+# Copyright 2018-2023 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 exeuction with default qubit 2 independent of any interface."""
+import pytest
+
+import pennylane as qml
+from pennylane import numpy as np
+from pennylane.interfaces.execution import _batch_transform, _preprocess_expand_fn
+
+from pennylane.devices.experimental import DefaultQubit2
+
+
+class TestPreprocessExpandFn:
+    """Tests the _preprocess_expand_fn helper function."""
+
+    def test_provided_is_callable(self):
+        """Test that if the expand_fn is not "device", it is simply returned."""
+
+        dev = DefaultQubit2()
+
+        def f(tape):
+            return tape
+
+        out = _preprocess_expand_fn(f, dev, 10)
+        assert out is f
+
+    def test_new_device_blank_expand_fn(self):
+        """Test that the expand_fn is blank if is new device."""
+
+        dev = DefaultQubit2()
+
+        out = _preprocess_expand_fn("device", dev, 10)
+
+        x = [1]
+        assert out(x) is x
+
+
+class TestBatchTransformHelper:
+    """Unit tests for the _batch_transform helper function."""
+
+    def test_warns_if_requested_off(self):
+        """Test that a warning is raised if the the batch transform is requested to not be used."""
+
+        # pylint: disable=too-few-public-methods
+        class CustomOp(qml.operation.Operator):
+            """Dummy operator."""
+
+            def decomposition(self):
+                return [qml.PauliX(self.wires[0])]
+
+        dev = DefaultQubit2()
+
+        qs = qml.tape.QuantumScript([CustomOp(0)], [qml.expval(qml.PauliZ(0))])
+
+        config = qml.devices.experimental.ExecutionConfig()
+
+        with pytest.warns(UserWarning, match="device batch transforms cannot be turned off"):
+            new_batch, post_procesing_fn, _ = _batch_transform(
+                (qs, qs), dev, config, device_batch_transform=False
+            )
+
+        assert len(new_batch) == 2
+
+        assert post_procesing_fn((1, 1)) == [1, 1]
+
+    def test_split_and_expand_performed(self):
+        """Test that preprocess returns the correct tapes when splitting and expanding
+        is needed."""
+
+        class NoMatOp(qml.operation.Operation):
+            """Dummy operation for expanding circuit."""
+
+            # pylint: disable=missing-function-docstring
+            num_wires = 1
+
+            # pylint: disable=arguments-renamed, invalid-overridden-method
+            @property
+            def has_matrix(self):
+                return False
+
+            def decomposition(self):
+                return [qml.PauliX(self.wires), qml.PauliY(self.wires)]
+
+        ops = [qml.Hadamard(0), NoMatOp(1), qml.RX([np.pi, np.pi / 2], wires=1)]
+        # Need to specify grouping type to transform tape
+        measurements = [qml.expval(qml.PauliX(0)), qml.expval(qml.PauliZ(1))]
+        tapes = [
+            qml.tape.QuantumScript(ops=ops, measurements=[measurements[0]]),
+            qml.tape.QuantumScript(ops=ops, measurements=[measurements[1]]),
+        ]
+
+        dev = DefaultQubit2()
+        config = qml.devices.experimental.ExecutionConfig(gradient_method="adjoint")
+
+        res_tapes, batch_fn, new_config = _batch_transform(tapes, dev, config)
+        expected_ops = [
+            [qml.Hadamard(0), qml.PauliX(1), qml.PauliY(1), qml.RX(np.pi, wires=1)],
+            [qml.Hadamard(0), qml.PauliX(1), qml.PauliY(1), qml.RX(np.pi / 2, wires=1)],
+        ]
+
+        assert len(res_tapes) == 4
+        for i, t in enumerate(res_tapes):
+            for op, expected_op in zip(t.operations, expected_ops[i % 2]):
+                assert qml.equal(op, expected_op)
+            assert len(t.measurements) == 1
+            if i < 2:
+                assert qml.equal(t.measurements[0], measurements[0])
+            else:
+                assert qml.equal(t.measurements[0], measurements[1])
+
+        input = ([[1, 2]], [[3, 4]], [[5, 6]], [[7, 8]])
+        assert np.array_equal(batch_fn(input), np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]))
+
+        assert new_config.grad_on_execution
+        assert new_config.use_device_gradient
+
+
+class TestNewDeviceIntegration:
+    """Localized tests for specific warnings, errors, and edge behaviour."""
+
+    def test_warning_if_not_device_batch_transform(self):
+        """Test that a warning is raised if the users requests to not run device batch transform."""
+
+        # pylint: disable=too-few-public-methods
+        class CustomOp(qml.operation.Operator):
+            """Dummy operator."""
+
+            def decomposition(self):
+                return [qml.PauliX(self.wires[0])]
+
+        dev = DefaultQubit2()
+
+        qs = qml.tape.QuantumScript([CustomOp(0)], [qml.expval(qml.PauliZ(0))])
+
+        with pytest.warns(UserWarning, match="device batch transforms cannot be turned off"):
+            results = qml.execute([qs], dev, device_batch_transform=False)
+
+        assert len(results) == 1
+        assert qml.math.allclose(results[0], -1)
+
+    def test_error_if_return_types_not_enabled(self):
+        """Check that an error is raised if return types is not enabled."""
+        qml.disable_return()
+
+        dev = DefaultQubit2()
+
+        qs = qml.tape.QuantumScript([], [qml.state()])
+        with pytest.raises(
+            ValueError, match="New device interface only works with return types enabled"
+        ):
+            qml.execute([qs], dev)
+
+        qml.enable_return()
+
+
+@pytest.mark.parametrize("gradient_fn", (None, "backprop", qml.gradients.param_shift))
+def test_caching(gradient_fn):
+    """Test that cache execute returns the cached result if the same script is executed
+    multiple times, both in multiple times in a batch and in separate batches."""
+    dev = DefaultQubit2()
+
+    qs = qml.tape.QuantumScript([qml.PauliX(0)], [qml.expval(qml.PauliZ(0))])
+
+    cache = {}
+
+    with qml.Tracker(dev) as tracker:
+        results = qml.execute([qs, qs], dev, cache=cache, gradient_fn=gradient_fn)
+        results2 = qml.execute([qs, qs], dev, cache=cache, gradient_fn=gradient_fn)
+
+    assert len(cache) == 1
+    assert cache[qs.hash] == -1.0
+
+    assert results == [-1.0, -1.0]
+    assert results2 == [-1.0, -1.0]
+
+    assert tracker.totals["batches"] == 1
+    assert tracker.totals["executions"] == 1
+    assert cache[qs.hash] == -1.0
diff --git a/tests/interfaces/test_autograd_new.py b/tests/interfaces/test_autograd_new.py
index c960f716e..683045098 100644
--- a/tests/interfaces/test_autograd_new.py
+++ b/tests/interfaces/test_autograd_new.py
@@ -118,7 +118,7 @@ class TestAutogradExecuteUnitTests:
             tape = qml.tape.QuantumScript.from_queue(q)
             return qml.execute([tape], device, gradient_fn=param_shift, interface="None")[0]
 
-        with pytest.raises(ValueError, match="Unknown interface"):
+        with pytest.raises(ValueError, match="interface must be in"):
             cost(a, device=dev)
 
     def test_grad_on_execution(self, mocker):
@@ -921,7 +921,6 @@ class TestHigherOrderDerivatives:
         """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.queuing.AnnotatedQueue() as q1:
diff --git a/tests/tests_passing_pylint.txt b/tests/tests_passing_pylint.txt
index aae5acb95..b2610d2e9 100644
--- a/tests/tests_passing_pylint.txt
+++ b/tests/tests_passing_pylint.txt
@@ -11,9 +11,11 @@ tests/drawer/test_draw_mpl.py
 tests/drawer/test_drawer_utils.py
 tests/drawer/test_tape_mpl.py
 tests/fourier/test_coefficients.py
+
 tests/gradients/core/**
 tests/gradients/finite_diff/**
 tests/gradients/parameter_shift/**
+tests/interfaces/default_qubit_2_integration/**
 tests/kernels/**
 tests/legacy/test_measurements_old.py
 tests/math/test_functions.py