Repository: PennyLaneAI/pennylane
Pull request #4935 (2023-12-11 14:36:47)

Jax-jit can use device vjp
[sc-50418]

This PR binds device-provided VJPs to the JAX-JIT interface. It also switches the jax-jit interface to the "Jacobian Product Calculator" system of defining how to take Jacobians/ JVPs/ VJPs.

Since we need to place a pure-callback around the jacobian calculation, but we want to calculate results and jacobians at the same time when feasible, this PR adds the `JacobianProductCalculator.execute_and_compute_jacobian` method. This has the benefit over calling `execute_fn` and `compute_jacobian` separately, as it is always `grad_on_execution=True` when actually taking derivatives, and always `grad_on_execution=False` when purely calculating results.

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.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.
albi3ro: Ran some performance numbers 🚀 

<img width="1136" alt="Screenshot 2023-12-14 at 9 50 52 AM" src="https://github.com/PennyLaneAI/pennylane/assets/6364575/f4bca295-c27b-4958-80d2-3f42bcfecdca">

Shows similar behaviour to the timings for torch and autograd. 

Interestingly enough, we have support for jacobians with jit, even though we don't with non-jit.  I think this is because we aren't setting `vectorized=True` in the pure-callback around the vjp calculation.  Something to revisit once the device derivative supports a broadcast dimension.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index f18c9d4f4..a4cfb3f3b 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -340,9 +340,10 @@
 
 <h4>Performance improvements and benchmarking</h4>
 
-* Autograd, PyTorch, and JAX (non-jit) can now use VJPs provided by the device from the new device API. If a device provides
+* Autograd, PyTorch, and JAX can now use VJPs provided by the device from the new device API. If a device provides
   a vector-Jacobian product, this can be selected by providing `device_vjp=True` to
   `qml.QNode` or `qml.execute`.
+  [(#4935)](https://github.com/PennyLaneAI/pennylane/pull/4935)
   [(#4557)](https://github.com/PennyLaneAI/pennylane/pull/4557)
   [(#4654)](https://github.com/PennyLaneAI/pennylane/pull/4654)
   [(#4878)](https://github.com/PennyLaneAI/pennylane/pull/4878)
diff --git a/pennylane/interfaces/autograd.py b/pennylane/interfaces/autograd.py
index 4c044b51d..0f3156f26 100644
--- a/pennylane/interfaces/autograd.py
+++ b/pennylane/interfaces/autograd.py
@@ -81,7 +81,7 @@ by the ``cache_full_jacobian`` keyword argument to :class:`~.TransformJacobianPr
 Other interfaces are capable of calculating the full jacobian in one call, so this patch is only present for autograd.
 
 """
-# pylint: disable=too-many-arguments
+# pylint: disable=too-many-arguments, unused-argument
 import logging
 from typing import Tuple, Callable
 
@@ -97,10 +97,12 @@ logger = logging.getLogger(__name__)
 logger.addHandler(logging.NullHandler())
 
 
+# pylint: disable=unused-argument
 def autograd_execute(
     tapes: Batch,
     execute_fn: ExecuteFn,
     jpc: qml.interfaces.jacobian_products.JacobianProductCalculator,
+    device=None,
 ):
     """Execute a batch of tapes with Autograd parameters on a device.
 
@@ -132,7 +134,6 @@ def autograd_execute(
     tapes = tuple(tapes)
     if logger.isEnabledFor(logging.DEBUG):
         logger.debug("Entry with (tapes=%s, execute_fn=%s, jpc=%s)", tapes, execute_fn, jpc)
-    # pylint: disable=unused-argument
     for tape in tapes:
         # set the trainable parameters
         params = tape.get_parameters(trainable_only=False)
diff --git a/pennylane/interfaces/execution.py b/pennylane/interfaces/execution.py
index ed99ee095..abdf2bb39 100644
--- a/pennylane/interfaces/execution.py
+++ b/pennylane/interfaces/execution.py
@@ -44,7 +44,7 @@ logger.addHandler(logging.NullHandler())
 
 device_type = Union[qml.Device, "qml.devices.Device"]
 
-jpc_interfaces = {"autograd", "numpy", "torch", "pytorch", "jax", "jax-python"}
+jpc_interfaces = {"autograd", "numpy", "torch", "pytorch", "jax", "jax-python", "jax-jit"}
 
 INTERFACE_MAP = {
     None: "Numpy",
@@ -131,7 +131,10 @@ def _get_ml_boundary_execute(
             from .torch import execute as ml_boundary
 
         elif interface == "jax-jit":
-            from .jax_jit import execute as ml_boundary
+            if device_vjp:
+                from .jax_jit import jax_jit_vjp_execute as ml_boundary
+            else:
+                from .jax_jit import jax_jit_jvp_execute as ml_boundary
         else:  # interface in {"jax", "jax-python", "JAX"}:
             if device_vjp:
                 from .jax import jax_vjp_execute as ml_boundary
@@ -567,7 +570,7 @@ def execute(
 
         interface = get_jax_interface_name(tapes)
         # Only need to calculate derivatives with jax when we know it will be executed later.
-        if interface == "jax":
+        if interface in {"jax", "jax-jit"}:
             grad_on_execution = grad_on_execution if isinstance(gradient_fn, Callable) else False
 
     if device_vjp and isinstance(device, qml.Device):
@@ -758,13 +761,6 @@ def execute(
 
             gradient_fn = device_gradient_fn
 
-            # Adjoint Jacobian with backward pass and jitting needs the original circuit output state which
-            # can not be reused from the device if `grad_on_execution is False`.
-            if interface == "jax-jit":
-                use_device_state = gradient_kwargs.get("use_device_state", None)
-                if use_device_state:
-                    gradient_kwargs["use_device_state"] = False
-
     elif grad_on_execution is True:
         # In "forward" mode, gradients are automatically handled
         # within execute_and_gradients, so providing a gradient_fn
@@ -784,12 +780,17 @@ def execute(
         )
         for _ in range(1, max_diff):
             ml_boundary_execute = _get_ml_boundary_execute(interface, _grad_on_execution)
-            execute_fn = partial(ml_boundary_execute, execute_fn=execute_fn, jpc=jpc)
+            execute_fn = partial(ml_boundary_execute, execute_fn=execute_fn, jpc=jpc, device=device)
             jpc = TransformJacobianProducts(execute_fn, gradient_fn, gradient_kwargs)
 
+            if interface == "jax-jit":
+                # no need to use pure callbacks around execute_fn or the jpc when taking
+                # higher order derivatives
+                interface = "jax"
+
     # trainable parameters can only be set on the first pass for jax
     # not higher order passes for higher order derivatives
-    if interface in {"jax", "jax-python"}:
+    if interface in {"jax", "jax-python", "jax-jit"}:
         for tape in tapes:
             params = tape.get_parameters(trainable_only=False)
             tape.trainable_params = qml.math.get_trainable_indices(params)
@@ -799,7 +800,7 @@ def execute(
     )
 
     if interface in jpc_interfaces:
-        results = ml_boundary_execute(tapes, execute_fn, jpc)
+        results = ml_boundary_execute(tapes, execute_fn, jpc, device=device)
     else:
         results = ml_boundary_execute(
             tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_diff=max_diff
diff --git a/pennylane/interfaces/jacobian_products.py b/pennylane/interfaces/jacobian_products.py
index 65bf49b9f..ab08c07b1 100644
--- a/pennylane/interfaces/jacobian_products.py
+++ b/pennylane/interfaces/jacobian_products.py
@@ -146,7 +146,7 @@ class JacobianProductCalculator(abc.ABC):
         >>> expected_vjp0 = 0.5 * -np.sin(0.1)
         >>> qml.math.allclose(vjps[0], expected_vjp0)
         True
-        >>> expected_jvp1 = 2.0 * -np.sin(0.2) + 3.0 * np.cos(0.2)
+        >>> expected_vjp1 = 2.0 * -np.sin(0.2) + 3.0 * np.cos(0.2)
         >>> qml.math.allclose(vjps[1], expected_vjp1)
         True
 
@@ -159,7 +159,7 @@ class JacobianProductCalculator(abc.ABC):
     def compute_jacobian(self, tapes: Batch) -> Tuple:
         """Compute the full Jacobian for a batch of tapes.
 
-        This method is required to compute Jacobians in the ``jax-jit`` interface
+        This method is required to compute Jacobians in the ``tensorflow`` interface
 
         Args:
             tapes (tuple[.QuantumScript]): the batch of tapes to take the derivatives of
@@ -179,6 +179,33 @@ class JacobianProductCalculator(abc.ABC):
 
         """
 
+    @abc.abstractmethod
+    def execute_and_compute_jacobian(self, tapes: Batch) -> Tuple:
+        """Compute the results and the full Jacobian for a batch of tapes.
+
+        This method is required to compute Jacobians in the ``jax-jit`` interface
+
+        Args:
+            tapes (tuple[.QuantumScript]): the batch of tapes to take the derivatives of
+
+        **Examples:**
+
+        For an instance of :class:`~.JacobianProductCalculator` ``jpc``, we have:
+
+        >>> tape0 = qml.tape.QuantumScript([qml.RX(0.1, wires=0)], [qml.expval(qml.PauliZ(0))])
+        >>> tape1 = qml.tape.QuantumScript([qml.RY(0.2, wires=0)], [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliX(0))])
+        >>> batch = (tape0, tape1)
+        >>> results, jacs = jpc.execute_and_compute_jacobian(batch)
+        >>> results
+        (0.9950041652780258, (0.9800665778412417, 0.19866933079506116))
+        >>> jacs
+        (array(-0.09983342), (array(-0.19866933), array(0.98006658)))
+
+        While this method could support non-scalar parameters in theory, no implementation currently supports
+        jacobians with non-scalar parameters.
+
+        """
+
 
 class TransformJacobianProducts(JacobianProductCalculator):
     """Compute VJPs, JVPs and Jacobians via a gradient transform :class:`~.TransformDispatcher`.
@@ -270,6 +297,24 @@ class TransformJacobianProducts(JacobianProductCalculator):
         vjp_results = self._inner_execute(tuple(vjp_tapes))
         return tuple(processing_fn(vjp_results))
 
+    def execute_and_compute_jacobian(self, tapes: Batch):
+        if logger.isEnabledFor(logging.DEBUG):  # pragma: no cover
+            logger.debug("execute_and_compute_jacobian called with %s", tapes)
+
+        num_result_tapes = len(tapes)
+
+        partial_gradient_fn = partial(self._gradient_transform, **self._gradient_kwargs)
+        jac_tapes, jac_postprocessing = qml.transforms.map_batch_transform(
+            partial_gradient_fn, tapes
+        )
+
+        full_batch = tapes + tuple(jac_tapes)
+        full_results = self._inner_execute(full_batch)
+        results = full_results[:num_result_tapes]
+        jac_results = full_results[num_result_tapes:]
+        jacs = jac_postprocessing(jac_results)
+        return tuple(results), tuple(jacs)
+
     def compute_jacobian(self, tapes: Batch):
         if logger.isEnabledFor(logging.DEBUG):  # pragma: no cover
             logger.debug("compute_jacobian called with %s", tapes)
@@ -469,30 +514,7 @@ class DeviceDerivatives(JacobianProductCalculator):
         jacobians with non-scalar parameters.
 
         """
-        if tapes not in self._results_cache and tapes not in self._jacs_cache:
-            results, jacs = self._dev_execute_and_compute_derivatives(tapes)
-            self._results_cache[tapes] = results
-            self._jacs_cache[tapes] = jacs
-        else:
-            if tapes in self._results_cache:
-                if logger.isEnabledFor(logging.DEBUG):  # pragma: no cover
-                    logger.debug("%s : Retrieving results from cache.", self)
-                results = self._results_cache[tapes]
-            else:
-                results = self._dev_execute(tapes)
-                self._results_cache[tapes] = results
-
-            if tapes in self._jacs_cache:
-                if logger.isEnabledFor(logging.DEBUG):  # pragma: no cover
-                    logger.debug("%s : Retrieving jacobian from cache.", self)
-                jacs = self._jacs_cache[tapes]
-            else:
-                # Here the jac was not cached but the results were. This can not happen because results are never
-                # cached alone (note that in the else clause above computing only results, jac must already be present)
-                raise NotImplementedError(
-                    "No path to cache results without caching jac. This branch should not occur."
-                )
-
+        results, jacs = self.execute_and_compute_jacobian(tapes)
         jvps = _compute_jvps(jacs, tangents, tapes)
         return results, jvps
 
@@ -584,6 +606,33 @@ class DeviceDerivatives(JacobianProductCalculator):
         self._jacs_cache[tapes] = jacs
         return jacs
 
+    def execute_and_compute_jacobian(self, tapes):
+        if tapes not in self._results_cache and tapes not in self._jacs_cache:
+            results, jacs = self._dev_execute_and_compute_derivatives(tapes)
+            self._results_cache[tapes] = results
+            self._jacs_cache[tapes] = jacs
+            return results, jacs
+
+        if tapes not in self._jacs_cache:
+            # Here the jac was not cached but the results were. This can not happen because results are never
+            # cached alone (note that in the else clause below computing only results, jac must already be present)
+            raise NotImplementedError(
+                "No path to cache results without caching jac. This branch should not occur."
+            )
+        if logger.isEnabledFor(logging.DEBUG):  # pragma: no cover
+            logger.debug("%s : Retrieving jacobian from cache.", self)
+        jacs = self._jacs_cache[tapes]
+
+        if tapes in self._results_cache:
+            if logger.isEnabledFor(logging.DEBUG):  # pragma: no cover
+                logger.debug("%s : Retrieving results from cache.", self)
+            results = self._results_cache[tapes]
+        else:
+            results = self._dev_execute(tapes)
+            self._results_cache[tapes] = results
+
+        return results, jacs
+
 
 class DeviceJacobianProducts(JacobianProductCalculator):
     """Compute jacobian products using the native device methods.
@@ -636,3 +685,9 @@ class DeviceJacobianProducts(JacobianProductCalculator):
             logger.debug("compute_jacobian called with %s", tapes)
         numpy_tapes = tuple(qml.transforms.convert_to_numpy_parameters(t) for t in tapes)
         return self._device.compute_derivatives(numpy_tapes, self._execution_config)
+
+    def execute_and_compute_jacobian(self, tapes: Batch) -> Tuple:
+        if logger.isEnabledFor(logging.DEBUG):  # pragma: no cover
+            logger.debug("execute_and_compute_jacobian called with %s", tapes)
+        numpy_tapes = tuple(qml.transforms.convert_to_numpy_parameters(t) for t in tapes)
+        return self._device.execute_and_compute_derivatives(numpy_tapes, self._execution_config)
diff --git a/pennylane/interfaces/jax.py b/pennylane/interfaces/jax.py
index c25e1533e..6079da9db 100644
--- a/pennylane/interfaces/jax.py
+++ b/pennylane/interfaces/jax.py
@@ -257,7 +257,7 @@ _execute_vjp = jax.custom_vjp(_execute_wrapper, nondiff_argnums=[1, 2, 3])
 _execute_vjp.defvjp(_vjp_fwd, _vjp_bwd)
 
 
-def jax_jvp_execute(tapes: Batch, execute_fn: ExecuteFn, jpc):
+def jax_jvp_execute(tapes: Batch, execute_fn: ExecuteFn, jpc, device=None):
     """Execute a batch of tapes with JAX parameters using JVP derivatives.
 
     Args:
@@ -279,7 +279,7 @@ def jax_jvp_execute(tapes: Batch, execute_fn: ExecuteFn, jpc):
     return _execute_jvp(parameters, _NonPytreeWrapper(tuple(tapes)), execute_fn, jpc)
 
 
-def jax_vjp_execute(tapes: Batch, execute_fn: ExecuteFn, jpc):
+def jax_vjp_execute(tapes: Batch, execute_fn: ExecuteFn, jpc, device=None):
     """Execute a batch of tapes with JAX parameters using VJP derivatives.
 
     Args:
diff --git a/pennylane/interfaces/jax_jit.py b/pennylane/interfaces/jax_jit.py
index d59761d76..5b064af71 100644
--- a/pennylane/interfaces/jax_jit.py
+++ b/pennylane/interfaces/jax_jit.py
@@ -12,62 +12,77 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """
-This module contains functions for adding the JAX interface
-to a PennyLane Device class.
+This module contains functions for binding JVPs or VJPs to JAX when using JIT.
+
+For information on registering VJPs and JVPs, please see the module documentation for ``jax.py``.
+
+When using JAX-JIT, we cannot convert arrays to numpy or act on their concrete values without
+using ``jax.pure_callback``.
+
+For example:
+
+>>> def f(x):
+...     return qml.math.unwrap(x)
+>>> x = jax.numpy.array(1.0)
+>>> jax.jit(f)(x)
+ValueError: Converting a JAX array to a NumPy array not supported when using the JAX JIT.
+>>> def g(x):
+...     expected_output_shape = jax.ShapeDtypeStruct((), jnp.float64)
+...     return jax.pure_callback(f, expected_output_shape, x)
+>>> jax.jit(g)(x)
+Array(1., dtype=float64)
+
+Note that we must provide the expected output shape for the function to use pure callbacks.
+
 """
-# pylint: disable=too-many-arguments, no-member
-from functools import partial
-from typing import Union, Tuple
+# pylint: disable=unused-argument, too-many-arguments
 import jax
 import jax.numpy as jnp
 
 import pennylane as qml
-from pennylane.transforms import convert_to_numpy_parameters
+from pennylane.typing import ResultBatch
 
 from .jacobian_products import _compute_jvps
 
+from .jax import _NonPytreeWrapper
+
 dtype = jnp.float64
 Zero = jax.custom_derivatives.SymbolicZero
 
 
-def _numeric_type_to_dtype(numeric_type):
-    """Auxiliary function for converting from Python numeric types to JAX
-    dtypes based on the precision defined for the interface."""
-
-    single_precision = dtype is jnp.float32
-    if numeric_type is int:
-        return jnp.int32 if single_precision else jnp.int64
-
-    if numeric_type is float:
-        return jnp.float32 if single_precision else jnp.float64
-
-    # numeric_type is complex
-    return jnp.complex64 if single_precision else jnp.complex128
-
+def _to_jax(result: qml.typing.ResultBatch) -> qml.typing.ResultBatch:
+    """Converts an arbitrary result batch to one with jax arrays.
+    Args:
+        result (ResultBatch): a nested structure of lists, tuples, and numpy arrays
+    Returns:
+        ResultBatch: a nested structure of tuples, and jax arrays
 
-def _set_copy_and_unwrap_tape(t, a, unwrap=True):
-    """Copy a given tape with operations and set parameters"""
-    tc = t.bind_new_parameters(a, list(range(len(a))))
-    return convert_to_numpy_parameters(tc) if unwrap else tc
+    """
+    # jax-jit not compatible with counts
+    # if isinstance(result, dict):
+    #     return result
+    if isinstance(result, (list, tuple)):
+        return tuple(_to_jax(r) for r in result)
+    return jnp.array(result)
 
 
-def set_parameters_on_copy_and_unwrap(tapes, params, unwrap=True):
+def _set_parameters_on_copy(tapes, params):
     """Copy a set of tapes with operations and set parameters"""
-    return tuple(_set_copy_and_unwrap_tape(t, a, unwrap=unwrap) for t, a in zip(tapes, params))
+    return tuple(t.bind_new_parameters(a, list(range(len(a)))) for t, a in zip(tapes, params))
 
 
-def _create_shape_dtype_struct(tape: "qml.tape.QuantumScript", device: "qml.Device"):
+def _result_shape_dtype_struct(tape: "qml.tape.QuantumScript", device: "qml.Device"):
     """Auxiliary function for creating the shape and dtype object structure
     given a tape."""
 
     shape = tape.shape(device)
     if len(tape.measurements) == 1:
-        tape_dtype = _numeric_type_to_dtype(tape.numeric_type)
+        tape_dtype = jnp.dtype(tape.numeric_type)
         if tape.shots.has_partitioned_shots:
             return tuple(jax.ShapeDtypeStruct(s, tape_dtype) for s in shape)
         return jax.ShapeDtypeStruct(tuple(shape), tape_dtype)
 
-    tape_dtype = tuple(_numeric_type_to_dtype(elem) for elem in tape.numeric_type)
+    tape_dtype = tuple(jnp.dtype(elem) for elem in tape.numeric_type)
     if tape.shots.has_partitioned_shots:
         return tuple(
             tuple(jax.ShapeDtypeStruct(tuple(s), d) for s, d in zip(si, tape_dtype)) for si in shape
@@ -91,7 +106,7 @@ def _jac_shape_dtype_struct(tape: "qml.tape.QuantumScript", device: "qml.Device"
     >>> fn(dev.execute(tapes))
     (array(0.), array([-0.42073549,  0.42073549]))
     """
-    shape_and_dtype = _create_shape_dtype_struct(tape, device)
+    shape_and_dtype = _result_shape_dtype_struct(tape, device)
     if len(tape.trainable_params) == 1:
         return shape_and_dtype
     if len(tape.measurements) == 1:
@@ -99,166 +114,106 @@ def _jac_shape_dtype_struct(tape: "qml.tape.QuantumScript", device: "qml.Device"
     return tuple(tuple(_s for _ in tape.trainable_params) for _s in shape_and_dtype)
 
 
-def _empty_jacs_for_shape(
-    shape_dtype: Union[Tuple, jax.ShapeDtypeStruct]
-) -> Union[Tuple, jnp.array]:
-    """Converts a nested tuple containing ``jax.ShapeDtypeStruct`` into a nested tuples
-    containing zeros arrays of the appropriate shape.
-
-    """
-    if isinstance(shape_dtype, tuple):
-        return tuple(_empty_jacs_for_shape(struct) for struct in shape_dtype)
-    return jnp.zeros(shape_dtype.shape, shape_dtype.dtype)
-
-
-def _switched_jacobian(tape, device, original_trainable_parameters, jacs):
-    """Adds in jacobians with of all zeros for parameters that had zero tangent.
+def _pytree_shape_dtype_struct(pytree):
+    """Creates a shape structure that matches the types and shapes for the provided pytree."""
+    leaves, struct = jax.tree_util.tree_flatten(pytree)
+    new_leaves = [jax.ShapeDtypeStruct(jnp.shape(l), l.dtype) for l in leaves]
+    return jax.tree_util.tree_unflatten(struct, new_leaves)
 
-    Used in ``execute_wrapper_jvp`` for grad_on_execution=True with device derivatives.
-
-    Note that this adds an additional nesting dimension to ``jacs``, with one jacobian
-    for each original trainable parameter. I am unsure why this works.
-
-    >>> dev = qml.devices.DefaultQubit()
-    >>> config = qml.devices.ExecutionConfig(gradient_method="adjoint")
-    >>> tape = qml.tape.QuantumTape([qml.RY(1.0, 0), qml.RX(0.6, 0), qml.RX(0.7, 0)], [qml.expval(qml.PauliZ(0))])
-    >>> tape.trainable_params = [0, 2]
-    >>> jac = dev.compute_derivatives(tape, config)
-    >>> jac
-    (array(-0.2250925), array(-0.52061271))
-    >>> _switched_jacobian(tape, dev, [0,1,2], jac)
-    ((array(-0.2250925), array(-0.52061271)),
-    (Array(0., dtype=float64), Array(0., dtype=float64)),
-    (array(-0.2250925), array(-0.52061271)))
 
+def _execute_wrapper(params, tapes, execute_fn, _, device) -> ResultBatch:
     """
-    intermediate_jacs = []
-
-    shape_dtype = _jac_shape_dtype_struct(tape, device)
-
-    for param_idx in original_trainable_parameters:
-        if param_idx in tape.trainable_params:
-            p_jac = jacs
-        else:
-            p_jac = _empty_jacs_for_shape(shape_dtype)
-        intermediate_jacs.append(p_jac)
-    return tuple(intermediate_jacs)
-
+    Execute tapes using a pure-callback.
+    """
+    shape_dtype_structs = tuple(_result_shape_dtype_struct(t, device) for t in tapes.vals)
+
+    def pure_callback_wrapper(p):
+        new_tapes = _set_parameters_on_copy(tapes.vals, p)
+        res = tuple(execute_fn(new_tapes))
+        # When executed under `jax.vmap` the `result_shapes_dtypes` will contain
+        # the shape without the vmap dimensions, while the function here will be
+        # executed with objects containing the vmap dimensions. So res[i].ndim
+        # will have an extra dimension for every `jax.vmap` wrapping this execution.
+        #
+        # The execute_fn will return an object with shape `(n_observables, batches)`
+        # but the default behaviour for `jax.pure_callback` is to add the extra
+        # dimension at the beginning, so `(batches, n_observables)`. So in here
+        # we detect with the heuristic above if we are executing under vmap and we
+        # swap the order in that case.
+        return jax.tree_map(lambda r, s: r.T if r.ndim > s.ndim else r, res, shape_dtype_structs)
+
+    out = jax.pure_callback(pure_callback_wrapper, shape_dtype_structs, params, vectorized=True)
+    return _to_jax(out)
+
+
+def _execute_and_compute_jvp(tapes, execute_fn, jpc, device, primals, tangents):
+    """
+    Compute the results and jvps using a pure callback around
+    :meth:`~.JacobianProductCalculator.execute_and_compute_jacobian`.
 
-# pylint: disable=unused-argument
-def _device_method_jac_via_callback(
-    params, tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n, max_diff
-):
-    """Perform a callback to compute the jacobian of tapes using a device method (e.g., adjoint).
+    Note that we must query the full jacobian inside the pure-callback so that jax can trace the JVP
+    calculation.
+    """
+    # Select the trainable params. Non-trainable params contribute a 0 gradient.
+    for tangent, tape in zip(tangents[0], tapes.vals):
+        tape.trainable_params = tuple(
+            idx for idx, t in enumerate(tangent) if not isinstance(t, Zero)
+        )
+    tangents_trainable = tuple(
+        tuple(t for t in tangent if not isinstance(t, Zero)) for tangent in tangents[0]
+    )
 
-    Note: we are not using the batch_jvp pipeline and rather split the steps of unwrapping tapes and the JVP
-    computation because:
+    def wrapper(inner_params):
+        new_tapes = _set_parameters_on_copy(tapes.vals, inner_params)
+        return jpc.execute_and_compute_jacobian(new_tapes)
 
-    1. Tape unwrapping has to happen in the callback: otherwise jitting is broken and Tracer objects
-    are converted to NumPy, something that raises an error;
+    res_struct = tuple(_result_shape_dtype_struct(t, device) for t in tapes.vals)
+    jac_struct = tuple(_jac_shape_dtype_struct(t, device) for t in tapes.vals)
+    results, jacobians = jax.pure_callback(wrapper, (res_struct, jac_struct), primals[0])
 
-    2. Passing in the tangents as an argument to the wrapper function called by the jax.pure_callback raises an
-    error (as of jax and jaxlib 0.3.25):
-    ValueError: Pure callbacks do not support transpose. Please use jax.custom_vjp to use callbacks while
-    taking gradients.
+    jvps = _compute_jvps(jacobians, tangents_trainable, tapes.vals)
 
-    Solution: Use the callback to compute the jacobian and then separately compute the JVP using the
-    tangent.
-    """
+    return _to_jax(results), _to_jax(jvps)
 
-    def wrapper(inner_params):
-        new_tapes = set_parameters_on_copy_and_unwrap(tapes, inner_params, unwrap=False)
-        return gradient_fn(new_tapes, **gradient_kwargs)
 
-    shape_dtype_structs = tuple(_jac_shape_dtype_struct(t, device) for t in tapes)
-    return jax.pure_callback(wrapper, shape_dtype_structs, params)
+def _vjp_fwd(params, tapes, execute_fn, jpc, device):
+    """Perform the forward pass execution, return results and the parameters as residuals."""
+    return _execute_wrapper(params, tapes, execute_fn, jpc, device), params
 
 
-# pylint: disable=unused-argument
-def _grad_transform_jac_via_callback(
-    params, tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n, max_diff
-):
-    """Perform a callback to compute the jacobian of tapes using a gradient transform (e.g., parameter-shift or
-    finite differences grad transform).
+def _vjp_bwd(tapes, execute_fn, jpc, device, params, dy):
+    """Perform the backward pass of a vjp calculation, returning the vjp."""
 
-    Note: we are not using the batch_jvp pipeline and rather split the steps of unwrapping tapes and the JVP
-    computation because:
+    def wrapper(inner_params, inner_dy):
+        new_tapes = _set_parameters_on_copy(tapes.vals, inner_params)
+        return tuple(jpc.compute_vjp(new_tapes, inner_dy))
 
-    1. Tape unwrapping has to happen in the callback: otherwise jitting is broken and Tracer objects
-    are converted to NumPy, something that raises an error;
+    vjp_shape = _pytree_shape_dtype_struct(params)
+    return (jax.pure_callback(wrapper, vjp_shape, params, dy),)
 
-    2. Passing in the tangents as an argument to the wrapper function called by the jax.pure_callback raises an
-    error (as of jax and jaxlib 0.3.25):
-    ValueError: Pure callbacks do not support transpose. Please use jax.custom_vjp to use callbacks while
-    taking gradients.
 
-    Solution: Use the callback to compute the jacobian and then separately compute the JVP using the
-    tangent.
-    """
+_execute_jvp_jit = jax.custom_jvp(_execute_wrapper, nondiff_argnums=[1, 2, 3, 4])
+_execute_jvp_jit.defjvp(_execute_and_compute_jvp, symbolic_zeros=True)
 
-    def wrapper(inner_params):
-        new_tapes = set_parameters_on_copy_and_unwrap(tapes, inner_params, unwrap=False)
-        all_jacs = []
-        for new_t in new_tapes:
-            jvp_tapes, res_processing_fn = gradient_fn(new_t, **gradient_kwargs)
-            jacs = execute_fn(jvp_tapes)[0]
-            jacs = res_processing_fn(jacs)
-            all_jacs.append(jacs)
-        return tuple(all_jacs)
-
-    expected_shapes = tuple(_jac_shape_dtype_struct(t, device) for t in tapes)
-    return jax.pure_callback(wrapper, expected_shapes, params)
-
-
-def _grad_transform_jac_no_callback(
-    params, tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n, max_diff
-):
-    """Performs the gradient transform in a differentiable manner."""
-    new_tapes = set_parameters_on_copy_and_unwrap(tapes, params, unwrap=False)
-    jacobians = []
-    for new_t in new_tapes:
-        jvp_tapes, res_processing_fn = gradient_fn(new_t, **gradient_kwargs)
-        jacs = execute(
-            jvp_tapes,
-            device,
-            execute_fn,
-            gradient_fn,
-            gradient_kwargs,
-            _n=_n + 1,
-            max_diff=max_diff,
-        )
-        jacs = res_processing_fn(jacs)
-        jacobians.append(jacs)
-    return tuple(jacobians)
+_execute_vjp_jit = jax.custom_vjp(_execute_wrapper, nondiff_argnums=[1, 2, 3, 4])
+_execute_vjp_jit.defvjp(_vjp_fwd, _vjp_bwd)
 
 
-def execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_diff=1):
-    """Execute a batch of tapes with JAX parameters on a device.
+def jax_jit_jvp_execute(tapes, execute_fn, jpc, device):
+    """Execute a batch of tapes with JAX parameters using VJP derivatives.
 
     Args:
         tapes (Sequence[.QuantumTape]): batch of tapes to execute
-        device (pennylane.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.
-        max_diff (int): If ``gradient_fn`` is a gradient transform, this option specifies
-            the maximum order of derivatives to support. Increasing this value allows
-            for higher order derivatives to be extracted, at the cost of additional
-            (classical) computational overhead during the backwards pass.
+        execute_fn (Callable[[Sequence[.QuantumTape]], ResultBatch]): a function that turns a batch of circuits into results
+        jpc (JacobianProductCalculator): a class that can compute the Jacobian for the input tapes.
+        device (pennylane.Device, pennylane.devices.Device): The device used for execution. Used to determine the shapes of outputs for
+            pure callback calls.
 
     Returns:
-        list[list[float]]: A nested list of tape results. Each element in
-        the returned list corresponds in order to the provided tapes.
+        TensorLike: A nested tuple of tape results. Each element in
+        the returned tuple corresponds in order to the provided tapes.
+
     """
-    # pylint: disable=unused-argument
 
     if any(
         m.return_type in (qml.measurements.Counts, qml.measurements.AllCounts)
@@ -269,170 +224,36 @@ def execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_d
         # not implemented and is required for the callback logic
         raise NotImplementedError("The JAX-JIT interface doesn't support qml.counts.")
 
-    parameters = tuple(list(t.get_parameters(trainable_only=False)) for t in tapes)
-
-    if gradient_fn is None:
-        return _execute_fwd(
-            parameters,
-            tapes=tapes,
-            device=device,
-            execute_fn=execute_fn,
-            gradient_kwargs=gradient_kwargs,
-            _n=_n,
-        )
-
-    return _execute_bwd(
-        parameters,
-        tapes=tapes,
-        device=device,
-        execute_fn=execute_fn,
-        gradient_fn=gradient_fn,
-        gradient_kwargs=gradient_kwargs,
-        _n=_n,
-        max_diff=max_diff,
-    )
-
-
-def _execute_bwd(
-    params,
-    tapes=None,
-    device=None,
-    execute_fn=None,
-    gradient_fn=None,
-    gradient_kwargs=None,
-    _n=1,
-    max_diff=2,
-):  # pylint: disable=dangerous-default-value,unused-argument
-    @jax.custom_jvp
-    def execute_wrapper(inner_params):
-        shape_dtype_structs = tuple(_create_shape_dtype_struct(t, device) for t in tapes)
-
-        def wrapper(p):
-            """Compute the forward pass."""
-            new_tapes = set_parameters_on_copy_and_unwrap(tapes, p, unwrap=False)
-            res, _ = execute_fn(new_tapes, **gradient_kwargs)
-            res = tuple(res)
-
-            # When executed under `jax.vmap` the `result_shapes_dtypes` will contain
-            # the shape without the vmap dimensions, while the function here will be
-            # executed with objects containing the vmap dimensions. So res[i].ndim
-            # will have an extra dimension for every `jax.vmap` wrapping this execution.
-            #
-            # The execute_fn will return an object with shape `(n_observables, batches)`
-            # but the default behaviour for `jax.pure_callback` is to add the extra
-            # dimension at the beginning, so `(batches, n_observables)`. So in here
-            # we detect with the heuristic above if we are executing under vmap and we
-            # swap the order in that case.
-            return jax.tree_map(
-                lambda r, s: r.T if r.ndim > s.ndim else r, res, shape_dtype_structs
-            )
-
-        return jax.pure_callback(wrapper, shape_dtype_structs, inner_params, vectorized=True)
-
-    # execute_wrapper_jvp merely registered, not used.
-    # pylint: disable=unused-variable
-    @partial(execute_wrapper.defjvp, symbolic_zeros=True)
-    def execute_wrapper_jvp(primals, tangents):
-        """Calculate the jvp in such a way that we can bind it to jax.
-
-        Closure Variables:
-            tapes, device, gradient_fn, gradient_kwargs, execute_fn, _n, max_diff
-        """
-        params = primals[0]
-
-        # Select the trainable params. Non-trainable params contribute a 0 gradient.
-        for tangent, tape in zip(tangents[0], tapes):
-            tape.trainable_params = tuple(
-                idx for idx, t in enumerate(tangent) if not isinstance(t, Zero)
-            )
-
-        if not isinstance(gradient_fn, qml.transforms.core.TransformDispatcher):
-            jacobians_func = _device_method_jac_via_callback
-        elif _n == max_diff:
-            jacobians_func = _grad_transform_jac_via_callback
-        else:
-            jacobians_func = _grad_transform_jac_no_callback
-
-        jacobians = jacobians_func(
-            params, tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n, max_diff
-        )
-
-        tangents_trainable = tuple(
-            tuple(t for t in tangent if not isinstance(t, Zero)) for tangent in tangents[0]
-        )
-        jvps = _compute_jvps(jacobians, tangents_trainable, tapes)
-
-        results = execute_wrapper(params)
-        return results, jvps
-
-    return execute_wrapper(params)
-
-
-# The execute function in forward mode
-def _execute_fwd(
-    params,
-    tapes=None,
-    device=None,
-    execute_fn=None,
-    gradient_kwargs=None,
-    _n=1,
-):  # pylint: disable=dangerous-default-value,unused-argument
-    """The auxiliary execute function for cases when the user requested
-    jacobians to be computed in forward mode (e.g. adjoint) or when no gradient function was
-    provided. This function does not allow multiple derivatives. It currently does not support shot vectors
-    because adjoint jacobian for default qubit does not support it."""
-
-    # pylint: disable=unused-variable
-    @jax.custom_jvp
-    def execute_wrapper(params):
-        def wrapper(p):
-            """Compute the forward pass."""
-            new_tapes = set_parameters_on_copy_and_unwrap(tapes, p, unwrap=False)
-            return execute_fn(new_tapes, **gradient_kwargs)
-
-        shape_dtype_structs = tuple(_create_shape_dtype_struct(t, device) for t in tapes)
-
-        jac_shape = tuple(_jac_shape_dtype_struct(t, device) for t in tapes)
-        jac_shape = jac_shape[0] if len(tapes) == 1 else jac_shape
+    parameters = tuple(tuple(t.get_parameters(trainable_only=False)) for t in tapes)
 
-        return jax.pure_callback(wrapper, (shape_dtype_structs, jac_shape), params)
+    return _execute_jvp_jit(parameters, _NonPytreeWrapper(tuple(tapes)), execute_fn, jpc, device)
 
-    @partial(execute_wrapper.defjvp, symbolic_zeros=True)
-    def execute_wrapper_jvp(primals, tangents):
-        """Primals[0] are parameters as Jax tracers and tangents[0] is a list of tangent vectors as Jax tracers.
 
-        Closure Variables:
-            tapes, execute_wrapper
-        """
+def jax_jit_vjp_execute(tapes, execute_fn, jpc, device=None):
+    """Execute a batch of tapes with JAX parameters using VJP derivatives.
 
-        # Get the original trainable parameters and the trainable parameters from symbolic zeros
-        original_trainable_parameters = [tape.trainable_params for tape in tapes]
-
-        for tangent, tape in zip(tangents[0], tapes):
-            tape.trainable_params = [
-                idx for idx, t in enumerate(tangent) if not isinstance(t, Zero)
-            ]
-
-        # Forward execution with the right trainable parameters
-        res, jacs = execute_wrapper(primals[0])
-
-        jacs_ = [jacs] if len(tapes) == 1 else jacs
-
-        updated_jacs = []
-
-        for tape, tape_jac, orig_trainable in zip(tapes, jacs_, original_trainable_parameters):
-            if tape.trainable_params != orig_trainable:
-                # Add zeros in the jacobians if the trainable params were switched
-                updated_jacs.append(_switched_jacobian(tape, device, orig_trainable, tape_jac))
-            else:
-                updated_jacs.append(tape_jac)
+    Args:
+        tapes (Sequence[.QuantumTape]): batch of tapes to execute
+        execute_fn (Callable[[Sequence[.QuantumTape]], ResultBatch]): a function that turns a batch of circuits into results
+        jpc (JacobianProductCalculator): a class that can compute the vector Jacobian product (VJP)
+            for the input tapes.
+        device (pennylane.Device, pennylane.devices.Device): The device used for execution. Used to determine the shapes of outputs for
+            pure callback calls.
 
-        updated_jacs = updated_jacs[0] if len(tapes) == 1 else tuple(updated_jacs)
+    Returns:
+        TensorLike: A nested tuple of tape results. Each element in
+        the returned tuple corresponds in order to the provided tapes.
 
-        # Get the jvps
-        tangents = [[t for t in tangent if not isinstance(t, Zero)] for tangent in tangents[0]]
-        jvps = _compute_jvps(jacs_, tangents, tapes)
+    """
+    if any(
+        m.return_type in (qml.measurements.Counts, qml.measurements.AllCounts)
+        for t in tapes
+        for m in t.measurements
+    ):
+        # Obtaining information about the shape of the Counts measurements is
+        # not implemented and is required for the callback logic
+        raise NotImplementedError("The JAX-JIT interface doesn't support qml.counts.")
 
-        return (res, updated_jacs), (jvps, updated_jacs)
+    parameters = tuple(tuple(t.get_parameters()) for t in tapes)
 
-    return execute_wrapper(params)[0]
+    return _execute_vjp_jit(parameters, _NonPytreeWrapper(tuple(tapes)), execute_fn, jpc, device)
diff --git a/pennylane/interfaces/torch.py b/pennylane/interfaces/torch.py
index 982c3150d..66366903c 100644
--- a/pennylane/interfaces/torch.py
+++ b/pennylane/interfaces/torch.py
@@ -60,7 +60,7 @@ result object.
 
 
 """
-# pylint: disable=too-many-arguments,protected-access,abstract-method
+# pylint: disable=too-many-arguments,protected-access,abstract-method,unused-argument
 import inspect
 import logging
 
@@ -189,7 +189,8 @@ class ExecuteTapes(torch.autograd.Function):
         return (None,) + vjps
 
 
-def execute(tapes, execute_fn, jpc):
+# pylint: disable=unused-argument
+def execute(tapes, execute_fn, jpc, device=None):
     """Execute a batch of tapes with Torch parameters on a device.
 
     Args:
@@ -211,7 +212,6 @@ def execute(tapes, execute_fn, jpc):
             jpc,
         )
 
-    # pylint: disable=unused-argument
     parameters = []
     for tape in tapes:
         # set the trainable parameters


Test patch:
diff --git a/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py
index 1b97a52a7..0f830f4de 100644
--- a/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py
@@ -16,19 +16,24 @@
 from functools import partial
 import pytest
 
+from param_shift_dev import ParamShiftDerivativesDevice
+
 import pennylane as qml
 from pennylane import numpy as np
 from pennylane import qnode
 from pennylane.devices import DefaultQubit
 
+
 qubit_device_and_diff_method = [
-    [DefaultQubit(), "backprop", True],
-    [DefaultQubit(), "finite-diff", False],
-    [DefaultQubit(), "parameter-shift", False],
-    [DefaultQubit(), "adjoint", True],
-    [DefaultQubit(), "adjoint", False],
-    [DefaultQubit(), "spsa", False],
-    [DefaultQubit(), "hadamard", False],
+    [DefaultQubit(), "backprop", True, False],
+    [DefaultQubit(), "finite-diff", False, False],
+    [DefaultQubit(), "parameter-shift", False, False],
+    [DefaultQubit(), "adjoint", True, False],
+    [DefaultQubit(), "adjoint", True, True],
+    [ParamShiftDerivativesDevice(), "device", False, True],
+    [DefaultQubit(), "adjoint", False, False],
+    [DefaultQubit(), "spsa", False, False],
+    [DefaultQubit(), "hadamard", False, False],
 ]
 interface_and_qubit_device_and_diff_method = [
     ["auto"] + inner_list for inner_list in qubit_device_and_diff_method
@@ -46,19 +51,26 @@ H_FOR_SPSA = 0.05
 
 
 @pytest.mark.parametrize(
-    "interface,dev,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method
+    "interface,dev,diff_method,grad_on_execution,device_vjp",
+    interface_and_qubit_device_and_diff_method,
 )
 class TestQNode:
     """Test that using the QNode with JAX integrates with the PennyLane
     stack"""
 
-    def test_execution_with_interface(self, dev, diff_method, grad_on_execution, interface):
+    def test_execution_with_interface(
+        self, dev, diff_method, grad_on_execution, interface, device_vjp
+    ):
         """Test execution works with the interface"""
-        if diff_method == "backprop":
-            pytest.skip("Test does not support backprop")
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention")
 
         @qnode(
-            dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface=interface,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a):
             qml.RY(a, wires=0)
@@ -78,7 +90,9 @@ class TestQNode:
         assert isinstance(grad, jax.Array)
         assert grad.shape == ()
 
-    def test_changing_trainability(self, dev, diff_method, grad_on_execution, interface, tol):
+    def test_changing_trainability(
+        self, dev, diff_method, grad_on_execution, interface, device_vjp, tol
+    ):
         """Test changing the trainability of parameters changes the
         number of differentiation requests made"""
         if diff_method != "parameter-shift":
@@ -92,6 +106,7 @@ class TestQNode:
             interface=interface,
             diff_method="parameter-shift",
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a, b):
             qml.RY(a, wires=0)
@@ -124,14 +139,21 @@ class TestQNode:
         circuit(a, b)
         assert circuit.qtape.trainable_params == [1]
 
-    def test_classical_processing(self, dev, diff_method, grad_on_execution, interface):
+    def test_classical_processing(self, dev, diff_method, grad_on_execution, device_vjp, interface):
         """Test classical processing within the quantum tape"""
         a = jax.numpy.array(0.1)
         b = jax.numpy.array(0.2)
         c = jax.numpy.array(0.3)
 
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention.")
+
         @qnode(
-            dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution
+            dev,
+            diff_method=diff_method,
+            interface=interface,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a, b, c):
             qml.RY(a * c, wires=0)
@@ -146,14 +168,24 @@ class TestQNode:
 
         assert len(res) == 2
 
-    def test_matrix_parameter(self, dev, diff_method, grad_on_execution, interface, tol):
+    def test_matrix_parameter(
+        self, dev, diff_method, grad_on_execution, device_vjp, interface, tol
+    ):
         """Test that the jax interface works correctly
         with a matrix parameter"""
+
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention.")
+
         U = jax.numpy.array([[0, 1], [1, 0]])
         a = jax.numpy.array(0.1)
 
         @qnode(
-            dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution
+            dev,
+            diff_method=diff_method,
+            interface=interface,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(U, a):
             qml.QubitUnitary(U, wires=0)
@@ -166,10 +198,15 @@ class TestQNode:
         if diff_method == "finite-diff":
             assert circuit.qtape.trainable_params == [1]
 
-    def test_differentiable_expand(self, dev, diff_method, grad_on_execution, interface, tol):
+    def test_differentiable_expand(
+        self, dev, diff_method, grad_on_execution, device_vjp, interface, tol
+    ):
         """Test that operation and nested tape expansion
         is differentiable"""
 
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention.")
+
         gradient_kwargs = {}
         if diff_method == "spsa":
             gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
@@ -193,6 +230,7 @@ class TestQNode:
             diff_method=diff_method,
             interface=interface,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
             **gradient_kwargs
         )
         def circuit(a, p):
@@ -219,7 +257,7 @@ class TestQNode:
         )
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
-    def test_jacobian_options(self, dev, diff_method, grad_on_execution, interface):
+    def test_jacobian_options(self, dev, diff_method, grad_on_execution, device_vjp, interface):
         """Test setting jacobian options"""
         if diff_method != "finite-diff":
             pytest.skip("Test only applies to finite diff.")
@@ -232,6 +270,7 @@ class TestQNode:
             h=1e-8,
             approx_order=2,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a):
             qml.RY(a[0], wires=0)
@@ -246,15 +285,19 @@ class TestQNode:
 
 
 @pytest.mark.parametrize(
-    "interface,dev,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method
+    "interface,dev,diff_method,grad_on_execution, device_vjp",
+    interface_and_qubit_device_and_diff_method,
 )
 class TestVectorValuedQNode:
     """Test that using vector-valued QNodes with JAX integrate with the
     PennyLane stack"""
 
-    def test_diff_expval_expval(self, dev, diff_method, grad_on_execution, interface, tol):
+    def test_diff_expval_expval(
+        self, dev, diff_method, grad_on_execution, device_vjp, interface, tol
+    ):
         """Test jacobian calculation"""
-
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention.")
         gradient_kwargs = {}
 
         if diff_method == "spsa":
@@ -270,6 +313,7 @@ class TestVectorValuedQNode:
             diff_method=diff_method,
             interface=interface,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
             **gradient_kwargs
         )
         def circuit(a, b):
@@ -311,9 +355,12 @@ class TestVectorValuedQNode:
         assert res[1][1].shape == ()
         assert np.allclose(res[1][1], expected[1][1], atol=tol, rtol=0)
 
-    def test_jacobian_no_evaluate(self, dev, diff_method, grad_on_execution, interface, tol):
+    def test_jacobian_no_evaluate(
+        self, dev, diff_method, grad_on_execution, device_vjp, interface, tol
+    ):
         """Test jacobian calculation when no prior circuit evaluation has been performed"""
-
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention.")
         gradient_kwargs = {}
 
         if diff_method == "spsa":
@@ -329,6 +376,7 @@ class TestVectorValuedQNode:
             diff_method=diff_method,
             interface=interface,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
             **gradient_kwargs
         )
         def circuit(a, b):
@@ -369,9 +417,13 @@ class TestVectorValuedQNode:
                 assert r.shape == ()
                 assert np.allclose(r, e, atol=tol, rtol=0)
 
-    def test_diff_single_probs(self, dev, diff_method, grad_on_execution, interface, tol):
+    def test_diff_single_probs(
+        self, dev, diff_method, grad_on_execution, device_vjp, interface, tol
+    ):
         """Tests correct output shape and evaluation for a tape
         with a single prob output"""
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention.")
         gradient_kwargs = {}
         if diff_method == "adjoint":
             pytest.skip("Adjoint does not support probs")
@@ -388,6 +440,7 @@ class TestVectorValuedQNode:
             diff_method=diff_method,
             interface=interface,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
             **gradient_kwargs
         )
         def circuit(x, y):
@@ -417,9 +470,13 @@ class TestVectorValuedQNode:
         assert np.allclose(res[0], expected.T[0], atol=tol, rtol=0)
         assert np.allclose(res[1], expected.T[1], atol=tol, rtol=0)
 
-    def test_diff_multi_probs(self, dev, diff_method, grad_on_execution, interface, tol):
+    def test_diff_multi_probs(
+        self, dev, diff_method, grad_on_execution, device_vjp, interface, tol
+    ):
         """Tests correct output shape and evaluation for a tape
         with multiple prob outputs"""
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention.")
         gradient_kwargs = {}
         if diff_method == "adjoint":
             pytest.skip("Adjoint does not support probs")
@@ -436,6 +493,7 @@ class TestVectorValuedQNode:
             diff_method=diff_method,
             interface=interface,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
             **gradient_kwargs
         )
         def circuit(x, y):
@@ -498,9 +556,13 @@ class TestVectorValuedQNode:
         assert jac[1][1].shape == (4,)
         assert np.allclose(jac[1][1], expected_1[1], atol=tol, rtol=0)
 
-    def test_diff_expval_probs(self, dev, diff_method, grad_on_execution, interface, tol):
+    def test_diff_expval_probs(
+        self, dev, diff_method, grad_on_execution, device_vjp, interface, tol
+    ):
         """Tests correct output shape and evaluation for a tape
         with prob and expval outputs"""
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention.")
         gradient_kwargs = {}
         if diff_method == "adjoint":
             pytest.skip("Adjoint does not support probs")
@@ -517,6 +579,7 @@ class TestVectorValuedQNode:
             diff_method=diff_method,
             interface=interface,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
             **gradient_kwargs
         )
         def circuit(x, y):
@@ -570,10 +633,12 @@ class TestVectorValuedQNode:
         assert np.allclose(jac[1][1], expected[1][1], atol=tol, rtol=0)
 
     def test_diff_expval_probs_sub_argnums(
-        self, dev, diff_method, grad_on_execution, interface, tol
+        self, dev, diff_method, grad_on_execution, device_vjp, interface, tol
     ):
         """Tests correct output shape and evaluation for a tape with prob and expval outputs with less
         trainable parameters (argnums) than parameters."""
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention.")
         kwargs = {}
         if diff_method == "adjoint":
             pytest.skip("Adjoint does not support probs")
@@ -589,6 +654,7 @@ class TestVectorValuedQNode:
             diff_method=diff_method,
             interface=interface,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
             **kwargs
         )
         def circuit(x, y):
@@ -621,9 +687,11 @@ class TestVectorValuedQNode:
         assert jac[1][0].shape == (2,)
         assert np.allclose(jac[1][0], expected[1][0], atol=tol, rtol=0)
 
-    def test_diff_var_probs(self, dev, diff_method, grad_on_execution, interface, tol):
+    def test_diff_var_probs(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol):
         """Tests correct output shape and evaluation for a tape
         with prob and variance outputs"""
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention.")
         gradient_kwargs = {}
         if diff_method == "adjoint":
             pytest.skip("Adjoint does not support probs")
@@ -642,6 +710,7 @@ class TestVectorValuedQNode:
             diff_method=diff_method,
             interface=interface,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
             **gradient_kwargs
         )
         def circuit(x, y):
@@ -819,12 +888,13 @@ class TestShotsIntegration:
 
 
 @pytest.mark.parametrize(
-    "interface,dev,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method
+    "interface,dev,diff_method,grad_on_execution, device_vjp",
+    interface_and_qubit_device_and_diff_method,
 )
 class TestQubitIntegration:
     """Tests that ensure various qubit circuits integrate correctly"""
 
-    def test_sampling(self, dev, diff_method, grad_on_execution, interface):
+    def test_sampling(self, dev, diff_method, grad_on_execution, device_vjp, interface):
         """Test sampling works as expected"""
         if grad_on_execution:
             pytest.skip("Sampling not possible with forward grad_on_execution differentiation.")
@@ -833,7 +903,11 @@ class TestQubitIntegration:
             pytest.skip("Adjoint warns with finite shots")
 
         @qnode(
-            dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution
+            dev,
+            diff_method=diff_method,
+            interface=interface,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit():
             qml.Hadamard(wires=[0])
@@ -849,7 +923,7 @@ class TestQubitIntegration:
         assert isinstance(res[1], jax.Array)
         assert res[1].shape == (10,)
 
-    def test_counts(self, dev, diff_method, grad_on_execution, interface):
+    def test_counts(self, dev, diff_method, grad_on_execution, device_vjp, interface):
         """Test counts works as expected"""
         if grad_on_execution:
             pytest.skip("Sampling not possible with forward grad_on_execution differentiation.")
@@ -858,7 +932,11 @@ class TestQubitIntegration:
             pytest.skip("Adjoint warns with finite shots")
 
         @qnode(
-            dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution
+            dev,
+            diff_method=diff_method,
+            interface=interface,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit():
             qml.Hadamard(wires=[0])
@@ -880,22 +958,32 @@ class TestQubitIntegration:
             assert isinstance(res[1], dict)
             assert len(res[1]) == 2
 
-    def test_chained_qnodes(self, dev, diff_method, grad_on_execution, interface):
+    def test_chained_qnodes(self, dev, diff_method, grad_on_execution, device_vjp, interface):
         """Test that the gradient of chained QNodes works without error"""
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention.")
 
         class Template(qml.templates.StronglyEntanglingLayers):
             def decomposition(self):
                 return [qml.templates.StronglyEntanglingLayers(*self.parameters, self.wires)]
 
         @qnode(
-            dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface=interface,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit1(weights):
             Template(weights, wires=[0, 1])
             return qml.expval(qml.PauliZ(0))
 
         @qnode(
-            dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface=interface,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit2(data, weights):
             qml.templates.AngleEmbedding(jax.numpy.stack([data, 0.7]), wires=[0, 1])
@@ -921,11 +1009,15 @@ class TestQubitIntegration:
 
         assert len(res) == 2
 
-    def test_postselection_differentiation(self, dev, diff_method, grad_on_execution, interface):
+    def test_postselection_differentiation(
+        self, dev, diff_method, grad_on_execution, device_vjp, interface
+    ):
         """Test that when postselecting with default.qubit, differentiation works correctly."""
 
         if diff_method in ["adjoint", "spsa", "hadamard"]:
             pytest.skip("Diff method does not support postselection.")
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention")
 
         @qml.qnode(
             dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution
@@ -938,7 +1030,11 @@ class TestQubitIntegration:
             return qml.expval(qml.PauliZ(1))
 
         @qml.qnode(
-            dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution
+            dev,
+            diff_method=diff_method,
+            interface=interface,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def expected_circuit(theta):
             qml.PauliX(1)
@@ -956,17 +1052,19 @@ class TestQubitIntegration:
 
 
 @pytest.mark.parametrize(
-    "interface,dev,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method
+    "interface,dev,diff_method,grad_on_execution, device_vjp",
+    interface_and_qubit_device_and_diff_method,
 )
 class TestQubitIntegrationHigherOrder:
     """Tests that ensure various qubit circuits integrate correctly when computing higher-order derivatives"""
 
-    def test_second_derivative(self, dev, diff_method, grad_on_execution, interface, tol):
+    def test_second_derivative(
+        self, dev, diff_method, grad_on_execution, device_vjp, interface, tol
+    ):
         """Test second derivative calculation of a scalar-valued QNode"""
-
         gradient_kwargs = {}
-        if diff_method == "adjoint":
-            pytest.skip("Adjoint does not second derivative.")
+        if diff_method in {"adjoint", "device"}:
+            pytest.skip("Adjoint does not support second derivatives.")
         elif diff_method == "spsa":
             gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
             gradient_kwargs["num_directions"] = 20
@@ -978,6 +1076,7 @@ class TestQubitIntegrationHigherOrder:
             diff_method=diff_method,
             interface=interface,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
             max_diff=2,
             **gradient_kwargs
         )
@@ -1008,10 +1107,10 @@ class TestQubitIntegrationHigherOrder:
         else:
             assert np.allclose(g2, expected_g2, atol=tol, rtol=0)
 
-    def test_hessian(self, dev, diff_method, grad_on_execution, interface, tol):
+    def test_hessian(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol):
         """Test hessian calculation of a scalar-valued QNode"""
         gradient_kwargs = {}
-        if diff_method == "adjoint":
+        if diff_method in {"adjoint", "device"}:
             pytest.skip("Adjoint does not support second derivative.")
         elif diff_method == "spsa":
             gradient_kwargs = {
@@ -1026,6 +1125,7 @@ class TestQubitIntegrationHigherOrder:
             diff_method=diff_method,
             interface=interface,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
             max_diff=2,
             **gradient_kwargs
         )
@@ -1059,10 +1159,12 @@ class TestQubitIntegrationHigherOrder:
         else:
             assert np.allclose(hess, expected_hess, atol=tol, rtol=0)
 
-    def test_hessian_vector_valued(self, dev, diff_method, grad_on_execution, interface, tol):
+    def test_hessian_vector_valued(
+        self, dev, diff_method, grad_on_execution, device_vjp, interface, tol
+    ):
         """Test hessian calculation of a vector-valued QNode"""
         gradient_kwargs = {}
-        if diff_method == "adjoint":
+        if diff_method in {"adjoint", "device"}:
             pytest.skip("Adjoint does not support second derivative.")
         elif diff_method == "spsa":
             gradient_kwargs = {
@@ -1077,6 +1179,7 @@ class TestQubitIntegrationHigherOrder:
             diff_method=diff_method,
             interface=interface,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
             max_diff=2,
             **gradient_kwargs
         )
@@ -1120,11 +1223,11 @@ class TestQubitIntegrationHigherOrder:
             assert np.allclose(hess, expected_hess, atol=tol, rtol=0)
 
     def test_hessian_vector_valued_postprocessing(
-        self, dev, diff_method, interface, grad_on_execution, tol
+        self, dev, diff_method, interface, device_vjp, grad_on_execution, tol
     ):
         """Test hessian calculation of a vector valued QNode with post-processing"""
         gradient_kwargs = {}
-        if diff_method == "adjoint":
+        if diff_method in {"adjoint", "device"}:
             pytest.skip("Adjoint does not support second derivative.")
         elif diff_method == "spsa":
             gradient_kwargs = {
@@ -1139,6 +1242,7 @@ class TestQubitIntegrationHigherOrder:
             diff_method=diff_method,
             interface=interface,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
             max_diff=2,
             **gradient_kwargs
         )
@@ -1185,11 +1289,11 @@ class TestQubitIntegrationHigherOrder:
             assert np.allclose(hess, expected_hess, atol=tol, rtol=0)
 
     def test_hessian_vector_valued_separate_args(
-        self, dev, diff_method, grad_on_execution, interface, tol
+        self, dev, diff_method, grad_on_execution, device_vjp, interface, tol
     ):
         """Test hessian calculation of a vector valued QNode that has separate input arguments"""
         gradient_kwargs = {}
-        if diff_method == "adjoint":
+        if diff_method in {"adjoint", "device"}:
             pytest.skip("Adjoint does not support second derivative.")
         elif diff_method == "spsa":
             gradient_kwargs = {
@@ -1205,6 +1309,7 @@ class TestQubitIntegrationHigherOrder:
             interface=interface,
             grad_on_execution=grad_on_execution,
             max_diff=2,
+            device_vjp=device_vjp,
             **gradient_kwargs
         )
         def circuit(a, b):
@@ -1249,7 +1354,7 @@ class TestQubitIntegrationHigherOrder:
         else:
             assert np.allclose(hess, expected_hess, atol=tol, rtol=0)
 
-    def test_state(self, dev, diff_method, grad_on_execution, interface, tol):
+    def test_state(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol):
         """Test that the state can be returned and differentiated"""
         if diff_method == "adjoint":
             pytest.skip("Adjoint does not support states")
@@ -1260,7 +1365,11 @@ class TestQubitIntegrationHigherOrder:
             dev._wires = qml.wires.Wires([0, 1])  # pylint:disable=protected-access
 
         @qnode(
-            dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution
+            dev,
+            diff_method=diff_method,
+            interface=interface,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x, y):
             qml.RX(x, wires=[0])
@@ -1284,9 +1393,13 @@ class TestQubitIntegrationHigherOrder:
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
     @pytest.mark.parametrize("state", [[1], [0, 1]])  # Basis state and state vector
-    def test_projector(self, state, dev, diff_method, grad_on_execution, interface, tol):
+    def test_projector(
+        self, state, dev, diff_method, grad_on_execution, device_vjp, interface, tol
+    ):
         """Test that the variance of a projector is correctly returned"""
         gradient_kwargs = {}
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention.")
         if diff_method == "adjoint":
             pytest.skip("Adjoint does not support projectors")
         elif diff_method == "hadamard":
@@ -1303,6 +1416,7 @@ class TestQubitIntegrationHigherOrder:
             diff_method=diff_method,
             interface=interface,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
             **gradient_kwargs
         )
         def circuit(x, y):
@@ -1326,7 +1440,8 @@ class TestQubitIntegrationHigherOrder:
 
 
 @pytest.mark.parametrize(
-    "interface,dev,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method
+    "interface,dev,diff_method,grad_on_execution, device_vjp",
+    interface_and_qubit_device_and_diff_method,
 )
 class TestTapeExpansion:
     """Test that tape expansion within the QNode integrates correctly
@@ -1334,7 +1449,7 @@ class TestTapeExpansion:
 
     @pytest.mark.parametrize("max_diff", [1, 2])
     def test_gradient_expansion_trainable_only(
-        self, dev, diff_method, grad_on_execution, max_diff, interface
+        self, dev, diff_method, grad_on_execution, device_vjp, max_diff, interface
     ):
         """Test that a *supported* operation with no gradient recipe is only
         expanded for parameter-shift and finite-differences when it is trainable."""
@@ -1353,6 +1468,7 @@ class TestTapeExpansion:
             grad_on_execution=grad_on_execution,
             max_diff=max_diff,
             interface=interface,
+            device_vjp=device_vjp,
         )
         def circuit(x, y):
             qml.Hadamard(wires=0)
@@ -1368,12 +1484,14 @@ class TestTapeExpansion:
 
     @pytest.mark.parametrize("max_diff", [1, 2])
     def test_hamiltonian_expansion_analytic(
-        self, dev, diff_method, grad_on_execution, max_diff, interface, mocker, tol
+        self, dev, diff_method, grad_on_execution, max_diff, device_vjp, interface, mocker, tol
     ):
         """Test that the Hamiltonian is not expanded if there
         are non-commuting groups and the number of shots is None
         and the first and second order gradients are correctly evaluated"""
         gradient_kwargs = {}
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradients transforms have a different vjp shape convention.")
         if diff_method == "adjoint":
             pytest.skip("The adjoint method does not yet support Hamiltonians")
         elif diff_method == "hadamard":
@@ -1395,6 +1513,7 @@ class TestTapeExpansion:
             diff_method=diff_method,
             grad_on_execution=grad_on_execution,
             max_diff=max_diff,
+            device_vjp=device_vjp,
             **gradient_kwargs
         )
         def circuit(data, weights, coeffs):
@@ -1439,13 +1558,15 @@ class TestTapeExpansion:
 
     @pytest.mark.parametrize("max_diff", [1, 2])
     def test_hamiltonian_finite_shots(
-        self, dev, diff_method, grad_on_execution, interface, max_diff, mocker
+        self, dev, diff_method, grad_on_execution, device_vjp, interface, max_diff, mocker
     ):
         """Test that the Hamiltonian is correctly measured if there
         are non-commuting groups and the number of shots is finite
         and the first and second order gradients are correctly evaluated"""
         gradient_kwargs = {}
         tol = 0.3
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention.")
         if diff_method in ("adjoint", "backprop", "finite-diff"):
             pytest.skip("The adjoint and backprop methods do not yet support sampling")
         elif diff_method == "hadamard":
@@ -1463,6 +1584,7 @@ class TestTapeExpansion:
             diff_method=diff_method,
             grad_on_execution=grad_on_execution,
             max_diff=max_diff,
+            device_vjp=device_vjp,
             **gradient_kwargs
         )
         def circuit(data, weights, coeffs):
@@ -1509,7 +1631,7 @@ class TestTapeExpansion:
     #         assert np.allclose(grad2_w_c, expected, atol=tol)
 
     def test_vmap_compared_param_broadcasting(
-        self, dev, diff_method, grad_on_execution, interface, tol
+        self, dev, diff_method, grad_on_execution, device_vjp, interface, tol
     ):
         """Test that jax.vmap works just as well as parameter-broadcasting with JAX JIT on the forward pass when
         vectorized=True is specified for the callback when caching is disabled."""
@@ -1533,6 +1655,7 @@ class TestTapeExpansion:
                 interface=interface,
                 diff_method=diff_method,
                 grad_on_execution=grad_on_execution,
+                device_vjp=device_vjp,
                 cache=None,
             )
             def _measure_operator():
@@ -1549,7 +1672,7 @@ class TestTapeExpansion:
         )
 
     def test_vmap_compared_param_broadcasting_multi_output(
-        self, dev, diff_method, grad_on_execution, interface, tol
+        self, dev, diff_method, grad_on_execution, device_vjp, interface, tol
     ):
         """Test that jax.vmap works just as well as parameter-broadcasting with JAX JIT on the forward pass when
         vectorized=True is specified for the callback when caching is disabled and when multiple output values
@@ -1574,6 +1697,7 @@ class TestTapeExpansion:
                 interface=interface,
                 diff_method=diff_method,
                 grad_on_execution=grad_on_execution,
+                device_vjp=device_vjp,
                 cache=None,
             )
             def _measure_operator():
@@ -1597,17 +1721,22 @@ jacobian_fn = [jax.jacobian, jax.jacrev, jax.jacfwd]
 
 @pytest.mark.parametrize("jacobian", jacobian_fn)
 @pytest.mark.parametrize(
-    "interface,dev,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method
+    "interface,dev,diff_method,grad_on_execution, device_vjp",
+    interface_and_qubit_device_and_diff_method,
 )
 class TestJIT:
     """Test JAX JIT integration with the QNode and automatic resolution of the
     correct JAX interface variant."""
 
-    def test_gradient(self, dev, diff_method, grad_on_execution, jacobian, tol, interface):
+    def test_gradient(
+        self, dev, diff_method, grad_on_execution, device_vjp, jacobian, tol, interface
+    ):
         """Test derivative calculation of a scalar valued QNode"""
         gradient_kwargs = {}
-        if diff_method == "adjoint":
-            pytest.xfail(reason="The adjoint method is not using host-callback currently")
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention.")
+        if device_vjp and jacobian == jax.jacfwd:
+            pytest.skip("device vjps not compatible with forward diff.")
         elif diff_method == "spsa":
             gradient_kwargs = {"h": H_FOR_SPSA, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)}
             tol = TOL_FOR_SPSA
@@ -1617,6 +1746,7 @@ class TestJIT:
             diff_method=diff_method,
             interface=interface,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
             **gradient_kwargs
         )
         def circuit(x):
@@ -1640,7 +1770,9 @@ class TestJIT:
         "ignore:Requested adjoint differentiation to be computed with finite shots."
     )
     @pytest.mark.parametrize("shots", [10, 1000])
-    def test_hermitian(self, dev, diff_method, grad_on_execution, shots, jacobian, interface):
+    def test_hermitian(
+        self, dev, diff_method, grad_on_execution, device_vjp, shots, jacobian, interface
+    ):
         """Test that the jax device works with qml.Hermitian and jitting even
         when shots>0.
 
@@ -1657,7 +1789,11 @@ class TestJIT:
         projector = np.array(qml.matrix(qml.PauliZ(0) @ qml.PauliZ(1)))
 
         @qml.qnode(
-            dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface=interface,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circ(projector):
             return qml.expval(qml.Hermitian(projector, wires=range(2)))
@@ -1669,7 +1805,9 @@ class TestJIT:
         "ignore:Requested adjoint differentiation to be computed with finite shots."
     )
     @pytest.mark.parametrize("shots", [10, 1000])
-    def test_probs_obs_none(self, dev, diff_method, grad_on_execution, shots, jacobian, interface):
+    def test_probs_obs_none(
+        self, dev, diff_method, grad_on_execution, device_vjp, shots, jacobian, interface
+    ):
         """Test that the jax device works with qml.probs, a MeasurementProcess
         that has obs=None even when shots>0."""
         # pylint: disable=unused-argument
@@ -1677,7 +1815,11 @@ class TestJIT:
             pytest.skip("Backpropagation is unsupported if shots > 0.")
 
         @qml.qnode(
-            dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface=interface,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit():
             return qml.probs(wires=0)
@@ -1687,14 +1829,20 @@ class TestJIT:
     @pytest.mark.xfail(
         reason="Non-trainable parameters are not being correctly unwrapped by the interface"
     )
-    def test_gradient_subset(self, dev, diff_method, grad_on_execution, jacobian, tol, interface):
+    def test_gradient_subset(
+        self, dev, diff_method, grad_on_execution, device_vjp, jacobian, tol, interface
+    ):
         """Test derivative calculation of a scalar valued QNode with respect
         to a subset of arguments"""
         a = jax.numpy.array(0.1)
         b = jax.numpy.array(0.2)
 
         @qnode(
-            dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution
+            dev,
+            diff_method=diff_method,
+            interface=interface,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a, b, c):
             qml.RY(a, wires=0)
@@ -1711,13 +1859,15 @@ class TestJIT:
         assert np.allclose(g, expected_g, atol=tol, rtol=0)
 
     def test_gradient_scalar_cost_vector_valued_qnode(
-        self, dev, diff_method, grad_on_execution, jacobian, tol, interface
+        self, dev, diff_method, grad_on_execution, device_vjp, jacobian, tol, interface
     ):
         """Test derivative calculation of a scalar valued cost function that
         uses the output of a vector-valued QNode"""
         gradient_kwargs = {}
         if diff_method == "adjoint":
-            pytest.xfail(reason="The adjoint method is not using host-callback currently")
+            pytest.xfail(reason="The adjoint method does not support probs")
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention.")
         elif diff_method == "spsa":
             gradient_kwargs = {"h": H_FOR_SPSA, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)}
             tol = TOL_FOR_SPSA
@@ -1727,6 +1877,7 @@ class TestJIT:
             diff_method=diff_method,
             interface=interface,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
             **gradient_kwargs
         )
         def circuit(x, y):
@@ -1759,13 +1910,24 @@ class TestJIT:
         assert np.allclose(g0, expected_g[0][idx], atol=tol, rtol=0)
         assert np.allclose(g1, expected_g[1][idx], atol=tol, rtol=0)
 
-    def test_matrix_parameter(self, dev, diff_method, grad_on_execution, jacobian, tol, interface):
+    # pylint: disable=unused-argument
+    def test_matrix_parameter(
+        self, dev, diff_method, grad_on_execution, device_vjp, jacobian, tol, interface
+    ):
         """Test that the JAX-JIT interface works correctly with a matrix
         parameter"""
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention.")
+        if jacobian == jax.jacfwd and device_vjp:
+            pytest.skip("device vjps are not compatible with forward differentiation.")
 
         # pylint: disable=unused-argument
         @qml.qnode(
-            dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution
+            dev,
+            diff_method=diff_method,
+            interface=interface,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circ(p, U):
             qml.QubitUnitary(U, wires=0)
@@ -1777,7 +1939,7 @@ class TestJIT:
         res = jax.jit(circ)(p, U)
         assert np.allclose(res, -np.cos(p), atol=tol, rtol=0)
 
-        jac_fn = jax.jit(jax.grad(circ, argnums=0))
+        jac_fn = jax.jit(jacobian(circ, argnums=0))
         res = jac_fn(p, U)
         assert np.allclose(res, np.sin(p), atol=tol, rtol=0)
 
@@ -1785,20 +1947,29 @@ class TestJIT:
 @pytest.mark.parametrize("shots", [None, 10000])
 @pytest.mark.parametrize("jacobian", jacobian_fn)
 @pytest.mark.parametrize(
-    "interface,dev,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method
+    "interface,dev,diff_method,grad_on_execution, device_vjp",
+    interface_and_qubit_device_and_diff_method,
 )
 class TestReturn:
     """Class to test the shape of the Grad/Jacobian with different return types."""
 
     def test_grad_single_measurement_param(
-        self, dev, diff_method, grad_on_execution, jacobian, shots, interface
+        self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface
     ):
         """For one measurement and one param, the gradient is a float."""
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention")
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
+        if jacobian == jax.jacfwd and device_vjp:
+            pytest.skip("jacfwd is not compatible with device_vjp=True.")
 
         @qnode(
-            dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface=interface,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a):
             qml.RY(a, wires=0)
@@ -1813,14 +1984,22 @@ class TestReturn:
         assert grad.shape == ()
 
     def test_grad_single_measurement_multiple_param(
-        self, dev, diff_method, grad_on_execution, jacobian, shots, interface
+        self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface
     ):
         """For one measurement and multiple param, the gradient is a tuple of arrays."""
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention.")
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
+        if jacobian == jax.jacfwd and device_vjp:
+            pytest.skip("jacfwd is not compatible with device_vjp=True.")
 
         @qnode(
-            dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface=interface,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a, b):
             qml.RY(a, wires=0)
@@ -1840,14 +2019,22 @@ class TestReturn:
         assert grad[1].shape == ()
 
     def test_grad_single_measurement_multiple_param_array(
-        self, dev, diff_method, grad_on_execution, jacobian, shots, interface
+        self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface
     ):
         """For one measurement and multiple param as a single array params, the gradient is an array."""
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention.")
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
+        if jacobian == jax.jacfwd and device_vjp:
+            pytest.skip("jacfwd is not compatible with device_vjp=True.")
 
         @qnode(
-            dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface=interface,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a):
             qml.RY(a[0], wires=0)
@@ -1862,10 +2049,12 @@ class TestReturn:
         assert grad.shape == (2,)
 
     def test_jacobian_single_measurement_param_probs(
-        self, dev, diff_method, grad_on_execution, jacobian, shots, interface
+        self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface
     ):
         """For a multi dimensional measurement (probs), check that a single array is returned with the correct
         dimension"""
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention.")
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
 
@@ -1873,7 +2062,11 @@ class TestReturn:
             pytest.skip("Test does not supports adjoint because of probabilities.")
 
         @qnode(
-            dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface=interface,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a):
             qml.RY(a, wires=0)
@@ -1888,17 +2081,23 @@ class TestReturn:
         assert jac.shape == (4,)
 
     def test_jacobian_single_measurement_probs_multiple_param(
-        self, dev, diff_method, grad_on_execution, jacobian, shots, interface
+        self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface
     ):
         """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with
         the correct dimension"""
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention.")
         if diff_method == "adjoint":
             pytest.skip("Test does not supports adjoint because of probabilities.")
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
 
         @qnode(
-            dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface=interface,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a, b):
             qml.RY(a, wires=0)
@@ -1919,17 +2118,23 @@ class TestReturn:
         assert jac[1].shape == (4,)
 
     def test_jacobian_single_measurement_probs_multiple_param_single_array(
-        self, dev, diff_method, grad_on_execution, jacobian, shots, interface
+        self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface
     ):
         """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with
         the correct dimension"""
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention.")
         if diff_method == "adjoint":
             pytest.skip("Test does not supports adjoint because of probabilities.")
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
 
         @qnode(
-            dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface=interface,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a):
             qml.RY(a[0], wires=0)
@@ -1943,17 +2148,25 @@ class TestReturn:
         assert jac.shape == (4, 2)
 
     def test_jacobian_expval_expval_multiple_params(
-        self, dev, diff_method, grad_on_execution, jacobian, shots, interface
+        self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface
     ):
         """The jacobian of multiple measurements with multiple params return a tuple of arrays."""
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention.")
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
+        if jacobian == jax.jacfwd and device_vjp:
+            pytest.skip("jacfwd is not compatible with device_vjp=True.")
 
         par_0 = jax.numpy.array(0.1)
         par_1 = jax.numpy.array(0.2)
 
         @qnode(
-            dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface=interface,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x, y):
             qml.RX(x, wires=[0])
@@ -1982,14 +2195,22 @@ class TestReturn:
         assert jac[1][1].shape == ()
 
     def test_jacobian_expval_expval_multiple_params_array(
-        self, dev, diff_method, grad_on_execution, jacobian, shots, interface
+        self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface
     ):
         """The jacobian of multiple measurements with a multiple params array return a single array."""
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention.")
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
+        if jacobian == jax.jacfwd and device_vjp:
+            pytest.skip("jacfwd is not compatible with device_vjp=True.")
 
         @qnode(
-            dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface=interface,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a):
             qml.RY(a[0], wires=0)
@@ -2010,9 +2231,11 @@ class TestReturn:
         assert jac[1].shape == (2,)
 
     def test_jacobian_var_var_multiple_params(
-        self, dev, diff_method, grad_on_execution, jacobian, shots, interface
+        self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface
     ):
         """The jacobian of multiple measurements with multiple params return a tuple of arrays."""
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention.")
         if diff_method == "adjoint":
             pytest.skip("Test does not supports adjoint because of var.")
         elif diff_method == "hadamard":
@@ -2024,7 +2247,11 @@ class TestReturn:
         par_1 = jax.numpy.array(0.2)
 
         @qnode(
-            dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface=interface,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x, y):
             qml.RX(x, wires=[0])
@@ -2054,9 +2281,11 @@ class TestReturn:
         assert jac[1][1].shape == ()
 
     def test_jacobian_var_var_multiple_params_array(
-        self, dev, diff_method, grad_on_execution, jacobian, shots, interface
+        self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface
     ):
         """The jacobian of multiple measurements with a multiple params array return a single array."""
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention.")
         if diff_method == "adjoint":
             pytest.skip("Test does not supports adjoint because of var.")
         elif diff_method == "hadamard":
@@ -2065,7 +2294,11 @@ class TestReturn:
             pytest.skip("Test does not support finite shots and adjoint/backprop")
 
         @qnode(
-            dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface=interface,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a):
             qml.RY(a[0], wires=0)
@@ -2086,9 +2319,13 @@ class TestReturn:
         assert jac[1].shape == (2,)
 
     def test_jacobian_multiple_measurement_single_param(
-        self, dev, diff_method, grad_on_execution, jacobian, shots, interface
+        self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface
     ):
         """The jacobian of multiple measurements with a single params return an array."""
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention")
+        if device_vjp and jacobian == jax.jacfwd:
+            pytest.skip("device vjp not compatible with forward differentiation.")
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
 
@@ -2096,7 +2333,11 @@ class TestReturn:
             pytest.skip("Test does not supports adjoint because of probabilities.")
 
         @qnode(
-            dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface=interface,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a):
             qml.RY(a, wires=0)
@@ -2117,16 +2358,22 @@ class TestReturn:
         assert jac[1].shape == (4,)
 
     def test_jacobian_multiple_measurement_multiple_param(
-        self, dev, diff_method, grad_on_execution, jacobian, shots, interface
+        self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface
     ):
         """The jacobian of multiple measurements with a multiple params return a tuple of arrays."""
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention.")
         if diff_method == "adjoint":
             pytest.skip("Test does not supports adjoint because of probabilities.")
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
 
         @qnode(
-            dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface=interface,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a, b):
             qml.RY(a, wires=0)
@@ -2156,16 +2403,22 @@ class TestReturn:
         assert jac[1][1].shape == (4,)
 
     def test_jacobian_multiple_measurement_multiple_param_array(
-        self, dev, diff_method, grad_on_execution, jacobian, shots, interface
+        self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface
     ):
         """The jacobian of multiple measurements with a multiple params array return a single array."""
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transforms have a different vjp shape convention.")
         if diff_method == "adjoint":
             pytest.skip("Test does not supports adjoint because of probabilities.")
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
 
         @qnode(
-            dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface=interface,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a):
             qml.RY(a[0], wires=0)
@@ -2195,16 +2448,17 @@ hessian_fn = [
 
 @pytest.mark.parametrize("hessian", hessian_fn)
 @pytest.mark.parametrize(
-    "interface,dev,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method
+    "interface,dev,diff_method,grad_on_execution, device_vjp",
+    interface_and_qubit_device_and_diff_method,
 )
 class TestReturnHessian:
     """Class to test the shape of the Hessian with different return types."""
 
     def test_hessian_expval_multiple_params(
-        self, dev, diff_method, hessian, grad_on_execution, interface
+        self, dev, diff_method, hessian, device_vjp, grad_on_execution, interface
     ):
         """The hessian of single a measurement with multiple params return a tuple of arrays."""
-        if diff_method == "adjoint":
+        if diff_method in {"adjoint", "device"}:
             pytest.skip("Test does not supports adjoint because second order diff.")
 
         par_0 = jax.numpy.array(0.1)
@@ -2216,6 +2470,7 @@ class TestReturnHessian:
             diff_method=diff_method,
             max_diff=2,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x, y):
             qml.RX(x, wires=[0])
@@ -2243,11 +2498,11 @@ class TestReturnHessian:
         assert hess[1][1].shape == ()
 
     def test_hessian_expval_multiple_param_array(
-        self, dev, diff_method, hessian, grad_on_execution, interface
+        self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface
     ):
         """The hessian of single measurement with a multiple params array return a single array."""
 
-        if diff_method == "adjoint":
+        if diff_method in {"adjoint", "device"}:
             pytest.skip("Test does not supports adjoint because second order diff.")
 
         params = jax.numpy.array([0.1, 0.2], dtype=jax.numpy.float64)
@@ -2258,6 +2513,7 @@ class TestReturnHessian:
             diff_method=diff_method,
             max_diff=2,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x):
             qml.RX(x[0], wires=[0])
@@ -2271,10 +2527,10 @@ class TestReturnHessian:
         assert hess.shape == (2, 2)
 
     def test_hessian_var_multiple_params(
-        self, dev, diff_method, hessian, grad_on_execution, interface
+        self, dev, diff_method, hessian, device_vjp, grad_on_execution, interface
     ):
         """The hessian of single a measurement with multiple params return a tuple of arrays."""
-        if diff_method == "adjoint":
+        if diff_method in {"adjoint", "device"}:
             pytest.skip("Test does not supports adjoint because second order diff.")
         elif diff_method == "hadamard":
             pytest.skip("Test does not supports hadamard because of var.")
@@ -2288,6 +2544,7 @@ class TestReturnHessian:
             diff_method=diff_method,
             max_diff=2,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x, y):
             qml.RX(x, wires=[0])
@@ -2315,10 +2572,10 @@ class TestReturnHessian:
         assert hess[1][1].shape == ()
 
     def test_hessian_var_multiple_param_array(
-        self, dev, diff_method, hessian, grad_on_execution, interface
+        self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface
     ):
         """The hessian of single measurement with a multiple params array return a single array."""
-        if diff_method == "adjoint":
+        if diff_method in {"adjoint", "device"}:
             pytest.skip("Test does not supports adjoint because second order diff.")
         elif diff_method == "hadamard":
             pytest.skip("Test does not supports hadamard because of var.")
@@ -2331,6 +2588,7 @@ class TestReturnHessian:
             diff_method=diff_method,
             max_diff=2,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x):
             qml.RX(x[0], wires=[0])
@@ -2344,10 +2602,10 @@ class TestReturnHessian:
         assert hess.shape == (2, 2)
 
     def test_hessian_probs_expval_multiple_params(
-        self, dev, diff_method, hessian, grad_on_execution, interface
+        self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface
     ):
         """The hessian of multiple measurements with multiple params return a tuple of arrays."""
-        if diff_method == "adjoint":
+        if diff_method in {"adjoint", "device"}:
             pytest.skip("Test does not supports adjoint because second order diff.")
         elif diff_method == "hadamard":
             pytest.skip("Test does not supports hadamard because of non commuting obs.")
@@ -2361,6 +2619,7 @@ class TestReturnHessian:
             diff_method=diff_method,
             max_diff=2,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x, y):
             qml.RX(x, wires=[0])
@@ -2387,11 +2646,11 @@ class TestReturnHessian:
                 assert h_comp.shape == (2,)
 
     def test_hessian_probs_expval_multiple_param_array(
-        self, dev, diff_method, hessian, grad_on_execution, interface
+        self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface
     ):
         """The hessian of multiple measurements with a multiple param array return a single array."""
 
-        if diff_method == "adjoint":
+        if diff_method in {"adjoint", "device"}:
             pytest.skip("Test does not supports adjoint because second order diff.")
         elif diff_method == "hadamard":
             pytest.skip("Test does not supports hadamard because of non commuting obs.")
@@ -2404,6 +2663,7 @@ class TestReturnHessian:
             diff_method=diff_method,
             max_diff=2,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x):
             qml.RX(x[0], wires=[0])
@@ -2422,10 +2682,10 @@ class TestReturnHessian:
         assert hess[1].shape == (2, 2, 2)
 
     def test_hessian_probs_var_multiple_params(
-        self, dev, diff_method, hessian, grad_on_execution, interface
+        self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface
     ):
         """The hessian of multiple measurements with multiple params return a tuple of arrays."""
-        if diff_method == "adjoint":
+        if diff_method in {"adjoint", "device"}:
             pytest.skip("Test does not supports adjoint because second order diff.")
         elif diff_method == "hadamard":
             pytest.skip("Test does not supports hadamard because of var.")
@@ -2439,6 +2699,7 @@ class TestReturnHessian:
             diff_method=diff_method,
             max_diff=2,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x, y):
             qml.RX(x, wires=[0])
@@ -2465,10 +2726,10 @@ class TestReturnHessian:
                 assert h_comp.shape == (2,)
 
     def test_hessian_probs_var_multiple_param_array(
-        self, dev, diff_method, hessian, grad_on_execution, interface
+        self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface
     ):
         """The hessian of multiple measurements with a multiple param array return a single array."""
-        if diff_method == "adjoint":
+        if diff_method in {"adjoint", "device"}:
             pytest.skip("Test does not supports adjoint because second order diff.")
         elif diff_method == "hadamard":
             pytest.skip("Test does not supports hadamard because of var.")
@@ -2481,6 +2742,7 @@ class TestReturnHessian:
             diff_method=diff_method,
             max_diff=2,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x):
             qml.RX(x[0], wires=[0])
@@ -2528,7 +2790,8 @@ def test_jax_device_hessian_shots(hessian, diff_method):
 @pytest.mark.parametrize("argnums", [0, 1, [0, 1]])
 @pytest.mark.parametrize("jacobian", jacobian_fn)
 @pytest.mark.parametrize(
-    "interface,dev,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method
+    "interface,dev,diff_method,grad_on_execution, device_vjp",
+    interface_and_qubit_device_and_diff_method,
 )
 class TestSubsetArgnums:
     def test_single_measurement(
@@ -2537,6 +2800,7 @@ class TestSubsetArgnums:
         dev,
         diff_method,
         grad_on_execution,
+        device_vjp,
         jacobian,
         argnums,
         jit_inside,
@@ -2544,6 +2808,10 @@ class TestSubsetArgnums:
     ):
         """Test single measurement with different diff methods with argnums."""
         kwargs = {}
+        if jacobian == jax.jacfwd and device_vjp:
+            pytest.skip("jacfwd is not compatible with device_vjp=True.")
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transform have a different vjp shape convention.")
         if diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
             tol = TOL_FOR_SPSA
@@ -2554,6 +2822,7 @@ class TestSubsetArgnums:
             diff_method=diff_method,
             grad_on_execution=grad_on_execution,
             cache=False,
+            device_vjp=device_vjp,
             **kwargs
         )
         def circuit(a, b):
@@ -2586,12 +2855,18 @@ class TestSubsetArgnums:
         dev,
         diff_method,
         grad_on_execution,
+        device_vjp,
         jacobian,
         argnums,
         jit_inside,
         tol,
     ):
         """Test multiple measurements with different diff methods with argnums."""
+        if jacobian == jax.jacfwd and device_vjp:
+            pytest.skip("jacfwd is not compatible with device_vjp=True.")
+        if dev.name == "param_shift.qubit":
+            pytest.xfail("gradient transform have a different vjp shape convention.")
+
         kwargs = {}
         if diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
@@ -2602,6 +2877,7 @@ class TestSubsetArgnums:
             interface=interface,
             diff_method=diff_method,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
             **kwargs
         )
         def circuit(a, b):
diff --git a/tests/interfaces/test_jacobian_products.py b/tests/interfaces/test_jacobian_products.py
index 34407ef81..bde1e50b8 100644
--- a/tests/interfaces/test_jacobian_products.py
+++ b/tests/interfaces/test_jacobian_products.py
@@ -221,8 +221,20 @@ class TestJacobianProductResults:
         x = 1.62
         tape = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.PauliZ(0))], shots=shots)
         jac = jpc.compute_jacobian((tape,))
+
         assert qml.math.allclose(jac, -np.sin(x), atol=_tol_for_shots(shots))
 
+    def test_execute_jacobian_basic(self, jpc, shots):
+        """Test execute_and_compute_jacobian for a simple single input single output."""
+        if shots and not _accepts_finite_shots(jpc):
+            pytest.skip("jpc does not work with finite shots.")
+
+        x = 1.62
+        tape = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.PauliZ(0))], shots=shots)
+        results, jacs = jpc.execute_and_compute_jacobian((tape,))
+        assert qml.math.allclose(results[0], np.cos(x), atol=_tol_for_shots(shots))
+        assert qml.math.allclose(jacs, -np.sin(x), atol=_tol_for_shots(shots))
+
     def test_batch_execute_jvp(self, jpc, shots):
         """Test execute_and_compute_jvp on a batch with ragged observables and parameters.."""
         if shots and not _accepts_finite_shots(jpc):
@@ -349,6 +361,66 @@ class TestJacobianProductResults:
                 jacs[1][1][1], -np.sin(x) * np.cos(y), atol=_tol_for_shots(shots)
             )
 
+    def test_batch_execute_jacobian(self, jpc, shots):
+        """Test execute_and_compute_jacobian on a batch with ragged observables and parameters."""
+
+        if shots and not _accepts_finite_shots(jpc):
+            pytest.skip("jpc does not work with finite shots.")
+        if jpc is hadamard_grad_jpc and qml.measurements.Shots(shots).has_partitioned_shots:
+            pytest.skip(
+                "hadamard gradient does not work with partitioned shots and multiple measurements."
+            )
+
+        x = np.array(0.28)
+        y = np.array(1.62)
+        phi = np.array(0.6293)
+
+        tape1 = qml.tape.QuantumScript(
+            [qml.RX(x, 0), qml.RY(y, 1), qml.CNOT((0, 1))],
+            [qml.expval(qml.PauliX(1)), qml.expval(qml.PauliY(0))],
+            shots=shots,
+        )
+        tape2 = qml.tape.QuantumScript(
+            [qml.Hadamard(0), qml.IsingXX(phi, wires=(0, 1))],
+            [qml.expval(qml.PauliZ(1))],
+            shots=shots,
+        )
+
+        # note reversed order of tapes in this test
+        res, jacs = jpc.execute_and_compute_jacobian((tape2, tape1))
+
+        if tape1.shots.has_partitioned_shots:
+            for i in [0, 1]:
+                assert qml.math.allclose(res[1][i][0], np.sin(y), atol=_tol_for_shots(shots))
+                assert qml.math.allclose(
+                    res[1][i][1], -np.sin(x) * np.sin(y), atol=_tol_for_shots(shots)
+                )
+                assert qml.math.allclose(res[0][i], np.cos(phi), atol=_tol_for_shots(shots))
+
+                assert qml.math.allclose(jacs[1][i][0][0], 0, atol=_tol_for_shots(shots))
+                assert qml.math.allclose(jacs[1][i][0][1], np.cos(y), atol=_tol_for_shots(shots))
+                assert qml.math.allclose(
+                    jacs[1][i][1][0], -np.cos(x) * np.sin(y), atol=_tol_for_shots(shots)
+                )
+                assert qml.math.allclose(
+                    jacs[1][i][1][1], -np.sin(x) * np.cos(y), atol=_tol_for_shots(shots)
+                )
+                assert qml.math.allclose(jacs[0][i], -np.sin(phi), atol=_tol_for_shots(shots))
+        else:
+            assert qml.math.allclose(res[1][0], np.sin(y), atol=_tol_for_shots(shots))
+            assert qml.math.allclose(res[1][1], -np.sin(x) * np.sin(y), atol=_tol_for_shots(shots))
+            assert qml.math.allclose(res[0], np.cos(phi), atol=_tol_for_shots(shots))
+
+            assert qml.math.allclose(jacs[0], -np.sin(phi), atol=_tol_for_shots(shots))
+            assert qml.math.allclose(jacs[1][0][0], 0, atol=_tol_for_shots(shots))
+            assert qml.math.allclose(jacs[1][0][1], np.cos(y), atol=_tol_for_shots(shots))
+            assert qml.math.allclose(
+                jacs[1][1][0], -np.cos(x) * np.sin(y), atol=_tol_for_shots(shots)
+            )
+            assert qml.math.allclose(
+                jacs[1][1][1], -np.sin(x) * np.cos(y), atol=_tol_for_shots(shots)
+            )
+
 
 @pytest.mark.parametrize("jpc", dev_jpc_matrix)
 class TestCachingDeviceDerivatives:
@@ -438,6 +510,31 @@ class TestCachingDeviceDerivatives:
         assert jpc._device.tracker.totals.get("derivatives", 0) == 0
         assert jpc._device.tracker.totals.get("executions", 0) == 0
 
+    def test_cached_on_execute_and_compute_jacobian(self, jpc):
+        """Test that execute_and_compute_jacobians caches results and Jacobians if they are not precalculated."""
+        x = 1.5
+        tape1 = qml.tape.QuantumScript(
+            [qml.Hadamard(0), qml.IsingXX(x, wires=(0, 1))], [qml.expval(qml.PauliZ(1))]
+        )
+        batch = (tape1,)
+
+        with jpc._device.tracker:
+            res, jacs = jpc.execute_and_compute_jacobian(batch)
+
+        assert jpc._device.tracker.totals["execute_and_derivative_batches"] == 1
+
+        assert qml.math.allclose(res, np.cos(x))
+        assert qml.math.allclose(jacs, -np.sin(x))
+
+        assert jpc._results_cache[batch] is res
+        assert qml.math.allclose(jpc._jacs_cache[batch], (-np.sin(x)))
+
+        with jpc._device.tracker:
+            jpc.execute_and_compute_jacobian(batch)
+
+        assert jpc._device.tracker.totals.get("derivatives", 0) == 0
+        assert jpc._device.tracker.totals.get("executions", 0) == 0
+
     def test_cached_on_vjps(self, jpc):
         """test that only jacs are cached on calls to compute_vjp."""
 
@@ -476,7 +573,7 @@ class TestCachingDeviceDerivatives:
         jpc._results_cache[batch] = "value"
 
         with pytest.raises(NotImplementedError):
-            jpc.execute_and_compute_jvp(batch, tuple())
+            jpc.execute_and_compute_jacobian(batch)
 
 
 @pytest.mark.parametrize("jpc", transform_jpc_matrix + [device_ps_jacs])
@@ -520,6 +617,50 @@ class TestProbsTransformJacobians:
         assert qml.math.allclose(jvp[1][0], 0)
         assert qml.math.allclose(jvp[1][1], -0.6 * np.sin(phi))
 
+    def test_execute_jacobian_multi_params_multi_out(self, jpc):
+        """Test execute_and_compute_jacobian with multiple parameters and multiple outputs"""
+        x = 0.93
+        y = -0.83
+        ops = [qml.RY(y, 0), qml.RX(x, 0)]
+        measurements = [qml.probs(wires=0), qml.expval(qml.PauliZ(0))]
+        tape1 = qml.tape.QuantumScript(ops, measurements)
+
+        phi = 0.545
+        ops2 = [qml.Hadamard(0), qml.IsingXX(phi, wires=(0, 1))]
+        measurements2 = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))]
+        tape2 = qml.tape.QuantumScript(ops2, measurements2)
+
+        res, jac = jpc.execute_and_compute_jacobian((tape1, tape2))
+
+        expected_res00 = 0.5 * np.array([1 + np.cos(x) * np.cos(y), 1 - np.cos(x) * np.cos(y)])
+        assert qml.math.allclose(res[0][0], expected_res00)
+
+        expected_res01 = np.cos(x) * np.cos(y)
+        assert qml.math.allclose(res[0][1], expected_res01)
+
+        assert qml.math.allclose(res[1][0], 0)
+        assert qml.math.allclose(res[1][1], np.cos(phi))
+
+        # first tape, first measurement, first parameters (y)
+        expected = 0.5 * np.array([-np.cos(x) * np.sin(y), np.cos(x) * np.sin(y)])
+        assert qml.math.allclose(jac[0][0][0], expected)
+
+        # first tape, first measurement, second parameter (x)
+        expected = 0.5 * np.array([-np.sin(x) * np.cos(y), np.sin(x) * np.cos(y)])
+        assert qml.math.allclose(jac[0][0][1], expected)
+
+        # first tape, second measurement, first parameter(y)
+        expected = -np.cos(x) * np.sin(y)
+        assert qml.math.allclose(jac[0][1][0], expected)
+        # first tape, second measurement, second parameter (x)
+        expected = -np.sin(x) * np.cos(y)
+        assert qml.math.allclose(jac[0][1][1], expected)
+
+        # second tape, first measurement, only parameter
+        assert qml.math.allclose(jac[1][0], 0)
+        # second tape, second measurement, only parameter
+        assert qml.math.allclose(jac[1][1], -np.sin(phi))
+
     def test_vjp_multi_params_multi_out(self, jpc):
         """Test compute_vjp with multiple parameters and multiple outputs."""
 
diff --git a/tests/interfaces/test_jax_jit.py b/tests/interfaces/test_jax_jit.py
index e61cf8387..8fe2e14c5 100644
--- a/tests/interfaces/test_jax_jit.py
+++ b/tests/interfaces/test_jax_jit.py
@@ -135,15 +135,20 @@ class TestJaxExecuteUnitTests:
         a = jax.numpy.array([0.1, 0.2])
         jax.jit(cost)(a)
 
-        # adjoint method only performs a single device execution, but gets both result and gradient
+        # adjoint method only performs a single device execution
+        # gradients are not requested when we only want the results
         assert dev.num_executions == 1
+        spy.assert_not_called()
+
+        # when the jacobian is requested, we always calculate it at the same time as the results
+        jax.grad(jax.jit(cost))(a)
         spy.assert_called()
 
     def test_no_gradients_on_execution(self, mocker):
         """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway"""
         dev = qml.device("default.qubit.legacy", wires=1)
         spy_execute = mocker.spy(qml.devices.DefaultQubitLegacy, "batch_execute")
-        spy_gradients = mocker.spy(qml.devices.DefaultQubitLegacy, "gradients")
+        spy_gradients = mocker.spy(qml.devices.DefaultQubitLegacy, "execute_and_gradients")
 
         def cost(a):
             with qml.queuing.AnnotatedQueue() as q: