Repository: PennyLaneAI/pennylane
Pull request #4963 (2023-12-20 16:30:06)

Add support for batched cotangents with adjoint VJP
**Context:**
This PR adds batching support to `adjoint_vjp`, which allows us to enable `jax.jacobian` with device provided VJPs on `default.qubit`.

**Description of the Change:**
* Add branches to `adjoint_vjp` to support batched cotangents, both for expvals and state measurement.
* Add logic for handling inhomogeneous cotangent shapes. This happens when one of the expvals returned by the tape is not used for the final result, in which case the cotangents will be a mix between numpy arrays with some constant number of elements, and the scalar value 0.
* Remove VJP execution related functions from `interfaces/jax.py`. Now, when using `device_vjp=True`, we always execute with `jax_jit_vjp_execute` from `interfaces/jax_jit.py`. This is done because we do not convert to numpy parameters when jitting, and converting to numpy parameters with causes the resulting jacobian computed with `adjoint_vjp` to be incorrect. The VJP code removed from `jax.py` was only used when `device_vjp=True`. Otherwise, the JVP code is used. Thank you @albi3ro for catching this!
* Remove `pytest.skip` and `pytest.xfail` on tests using `jax.jacobian` with adjoint diff. There were also a few tests that were still being skipped when non-expval measurements were made with adjoint diff. I also stopped skipping those tests.
* Added light documentation about the expected shape of the cotangents depending on what we are trying to compute with them.

**Benefits:**
* `jax.jacobian` works with adjoint.
* JAX is strictly faster with `device_vjp=True`.
* `jax.jacobian` is now vectorized when jitting with `device_vjp=True`.

**Possible Drawbacks:**
None

**Related GitHub Issues:**



codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4963?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`b6b1f16`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/b6b1f164b33bb9919ead465007a05ea628d8d551?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.67% compared to head [(`af22870`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4963?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.66%.


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #4963      +/-   ##
==========================================
- Coverage   99.67%   99.66%   -0.02%     
==========================================
  Files         392      392              
  Lines       35556    35300     -256     
==========================================
- Hits        35442    35183     -259     
- Misses        114      117       +3     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4963?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

mudit2812: [sc-52171]
mudit2812: Fixed
<img width="737" alt="image" src="https://github.com/PennyLaneAI/pennylane/assets/18223836/e7b6ea2b-4c4f-4114-ad86-907faa372573">


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 62eef074d..0f5795017 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -4,6 +4,10 @@
 
 <h3>New features since last release</h3>
 
+* Adjoint device VJP's are now supported with `jax.jacobian`. `device_vjp=True` is
+  is now strictly faster for jax.
+  [(#4963)](https://github.com/PennyLaneAI/pennylane/pull/4963)
+
 <h3>Improvements 🛠</h3>
 
 * Improve the performance of circuit-cutting workloads with large numbers of generated tapes.
@@ -116,4 +120,5 @@ Christina Lee,
 Isaac De Vlugt,
 Korbinian Kottmann,
 Lee J. O'Riordan,
+Mudit Pandey,
 Matthew Silverman.
diff --git a/pennylane/devices/default_qubit.py b/pennylane/devices/default_qubit.py
index 92b11bc84..cde6499dc 100644
--- a/pennylane/devices/default_qubit.py
+++ b/pennylane/devices/default_qubit.py
@@ -770,6 +770,45 @@ class DefaultQubit(Device):
         cotangents: Tuple[Number],
         execution_config: ExecutionConfig = DefaultExecutionConfig,
     ):
+        r"""The vector jacobian product used in reverse-mode differentiation. ``DefaultQubit`` uses the
+        adjoint differentiation method to compute the VJP.
+
+        Args:
+            circuits (Union[QuantumTape, Sequence[QuantumTape]]): the circuit or batch of circuits
+            cotangents (Tuple[Number, Tuple[Number]]): Gradient-output vector. Must have shape matching the output shape of the
+                corresponding circuit. If the circuit has a single output, `cotangents` may be a single number, not an iterable
+                of numbers.
+            execution_config (ExecutionConfig): a datastructure with all additional information required for execution
+
+        Returns:
+            tensor-like: A numeric result of computing the vector jacobian product
+
+        **Definition of vjp:**
+
+        If we have a function with jacobian:
+
+        .. math::
+
+            \vec{y} = f(\vec{x}) \qquad J_{i,j} = \frac{\partial y_i}{\partial x_j}
+
+        The vector jacobian product is the inner product of the derivatives of the output ``y`` with the
+        Jacobian matrix. The derivatives of the output vector are sometimes called the **cotangents**.
+
+        .. math::
+
+            \text{d}x_i = \Sigma_{i} \text{d}y_i J_{i,j}
+
+        **Shape of cotangents:**
+
+        The value provided to ``cotangents`` should match the output of :meth:`~.execute`. For computing the full Jacobian,
+        the cotangents can be batched to vectorize the computation. In this case, the cotangents can have the following
+        shapes. ``batch_size`` below refers to the number of entries in the Jacobian:
+
+        * For a state measurement, the cotangents must have shape ``(batch_size, 2 ** n_wires)``
+        * For ``n`` expectation values, the cotangents must have shape ``(n, batch_size)``. If ``n = 1``,
+          then the shape must be ``(batch_size,)``.
+
+        """
         is_single_circuit = False
         if isinstance(circuits, QuantumScript):
             is_single_circuit = True
diff --git a/pennylane/devices/qubit/adjoint_jacobian.py b/pennylane/devices/qubit/adjoint_jacobian.py
index 683c0683a..aa3fcaf30 100644
--- a/pennylane/devices/qubit/adjoint_jacobian.py
+++ b/pennylane/devices/qubit/adjoint_jacobian.py
@@ -79,8 +79,6 @@ def adjoint_jacobian(tape: QuantumTape, state=None):
 
         The adjoint differentiation method has the following restrictions:
 
-        * Only expectation values are supported as measurements.
-
         * Cannot differentiate with respect to observables.
 
         * Cannot differentiate with respect to state-prep operations.
@@ -158,8 +156,6 @@ def adjoint_jvp(tape: QuantumTape, tangents: Tuple[Number], state=None):
 
         The adjoint differentiation method has the following restrictions:
 
-        * Only expectation values are supported as measurements.
-
         * Cannot differentiate with respect to observables.
 
         * Observable being measured must have a matrix.
@@ -220,6 +216,106 @@ def adjoint_jvp(tape: QuantumTape, tangents: Tuple[Number], state=None):
     return tuple(np.array(t) for t in tangents_out)
 
 
+def _get_vjp_bras(tape, cotangents, ket):
+    """Helper function for getting the bras for adjoint vjp, the batch size of the
+    cotangents, as well as a list of indices for which the cotangents are zero.
+
+    Args:
+        tape (QuantumTape): circuit that the function takes the gradient of.
+        tangents (Tuple[Number]): gradient vector for input parameters.
+        ket (TensorLike): the final state of the circuit.
+
+    Returns:
+        Tuple[TensorLike, int, List]: The return contains the following:
+            * Final bra for batch size ``None``, else array of bras
+            * Batch size. None if cotangents are not batched
+            * List containing batch indices that are zero. Empty for unbatched
+              cotangents
+    """
+
+    if isinstance(tape.measurements[0], qml.measurements.StateMP):
+        batched_cotangents = np.ndim(cotangents) == 2
+        batch_size = np.shape(cotangents)[0] if batched_cotangents else None
+        bras = np.conj(cotangents.reshape(-1, *ket.shape))
+        bras = bras if batched_cotangents else np.squeeze(bras)
+        return bras, batch_size, []
+
+    # If not state measurement, measurements are guaranteed to be expectation values
+
+    single_cotangent = len(tape.measurements) == 1
+
+    if not single_cotangent:
+        # Pad cotangents if shape is inhomogenous
+        # inner_shape will only be None if cotangents is a vector. We assume that for
+        # inhomogenous cotangents, all non-scalar values have the same shape.
+        inner_shape = next((np.shape(cot) for cot in cotangents if np.shape(cot) != ()), None)
+        if inner_shape is not None:
+            # Batched cotangents. Find scalar zeros and pad to make the shape of cotangents
+            # homogenous
+            new_cotangents = []
+
+            for i, c in enumerate(cotangents):
+                if np.shape(c) == () and np.allclose(c, 0.0):
+                    new_cotangents.append(np.zeros(inner_shape))
+                else:
+                    new_cotangents.append(c)
+
+            cotangents = new_cotangents
+
+    cotangents = np.array(cotangents)
+    if single_cotangent:
+        # Expand dimensions for cases when there is a single broadcasted cotangent
+        # so that the cotangent has 2 dimensions, which is expected in the rest of the
+        # function for batched cotangents. For unbatched cases, this will make a scalar
+        # cotangent a one-item array
+        cotangents = np.expand_dims(cotangents, 0)
+
+    # One dimension for number of expectation values, one dimension for batch size.
+    batched_cotangents = np.ndim(cotangents) == 2
+    batch_size = cotangents.shape[1] if batched_cotangents else None
+    if np.allclose(cotangents, 0.0):
+        return None, batch_size, []
+
+    new_obs, null_batch_indices = [], []
+
+    # Collect list of observables to use for the adjoint algorithm. These will be used
+    # to construct the initial bras
+    if batched_cotangents:
+        for i, cots in enumerate(cotangents.T):
+            new_cs, new_os = [], []
+            for c, o in zip(cots, tape.observables):
+                if not np.allclose(c, 0.0):
+                    new_cs.append(c)
+                    new_os.append(o)
+            if len(new_cs) == 0:
+                null_batch_indices.append(i)
+            else:
+                new_obs.append(qml.dot(new_cs, new_os))
+
+    else:
+        new_cs, new_os = [], []
+        for c, o in zip(cotangents, tape.observables):
+            if not np.allclose(c, 0.0):
+                new_cs.append(c)
+                new_os.append(o)
+
+        new_obs.append(qml.dot(new_cs, new_os))
+
+    # Create bra(s) by taking product of observable(s) with the final state
+    bras = np.empty((len(new_obs), *ket.shape), dtype=ket.dtype)
+
+    for kk, obs in enumerate(new_obs):
+        if obs.pauli_rep is not None:
+            flat_bra = obs.pauli_rep.dot(ket.flatten(), wire_order=list(range(tape.num_wires)))
+            bras[kk] = 2 * flat_bra.reshape(ket.shape)
+        else:
+            bras[kk] = 2 * apply_operation(obs, ket)
+
+    bras = bras if batched_cotangents else np.squeeze(bras)
+
+    return bras, batch_size, null_batch_indices
+
+
 def adjoint_vjp(tape: QuantumTape, cotangents: Tuple[Number], state=None):
     """The vector jacobian product used in reverse-mode differentiation.
 
@@ -233,15 +329,21 @@ def adjoint_vjp(tape: QuantumTape, cotangents: Tuple[Number], state=None):
 
         The adjoint differentiation method has the following restrictions:
 
-        * Only expectation values are supported as measurements.
-
         * Cannot differentiate with respect to observables.
 
         * Observable being measured must have a matrix.
 
     Args:
         tape (QuantumTape): circuit that the function takes the gradient of
-        cotangents (Tuple[Number]): gradient vector for output parameters
+        cotangents (Tuple[Number]): gradient vector for output parameters. For computing
+            the full Jacobian, the cotangents can be batched to vectorize the computation.
+            In this case, the cotangents can have the following shapes. ``batch_size``
+            below refers to the number of entries in the Jacobian:
+
+            * For a state measurement, cotangents must have shape ``(batch_size, 2 ** n_wires)``.
+            * For ``n`` expectation values, the cotangents must have shape ``(n, batch_size)``.
+              If ``n = 1``, then the shape must be ``(batch_size,)``.
+
         state (TensorLike): the final state of the circuit; if not provided,
             the final state will be computed by executing the tape
 
@@ -250,7 +352,7 @@ def adjoint_vjp(tape: QuantumTape, cotangents: Tuple[Number], state=None):
     """
     # See ``adjoint_jacobian.md`` to more information on the algorithm.
 
-    # Map wires if custom wire labels used
+    # Map wires if custom wire labels used)
     if set(tape.wires) != set(range(tape.num_wires)):
         wire_map = {w: i for i, w in enumerate(tape.wires)}
         tapes, fn = qml.map_wires(tape, wire_map)
@@ -258,34 +360,19 @@ def adjoint_vjp(tape: QuantumTape, cotangents: Tuple[Number], state=None):
 
     ket = state if state is not None else get_final_state(tape)[0]
 
+    bras, batch_size, null_batch_indices = _get_vjp_bras(tape, cotangents, ket)
+    if bras is None:
+        # Cotangents are zeros
+        if batch_size is None:
+            return tuple(0.0 for _ in tape.trainable_params)
+        return tuple(np.zeros((len(tape.trainable_params), batch_size)))
+
     if isinstance(tape.measurements[0], qml.measurements.StateMP):
-        try:
-            bra = np.conj(cotangents.reshape(ket.shape))
-        except ValueError as e:
-            raise NotImplementedError("adjoint_vjp does not yet support JAX Jacobians") from e
 
         def real_if_expval(val):
             return val
 
     else:
-        cotangents = (cotangents,) if qml.math.shape(cotangents) == tuple() else cotangents
-        new_cotangents, new_observables = [], []
-        for c, o in zip(cotangents, tape.observables):
-            if qml.math.size(c) > 1:
-                raise NotImplementedError(
-                    "adjoint_vjp does not yet support JAX Jacobians, as they use a broadcast dimension on the cotangent dy."
-                )
-            if not np.allclose(c, 0.0):
-                new_cotangents.append(c)
-                new_observables.append(o)
-        if len(new_cotangents) == 0:
-            return tuple(0.0 for _ in tape.trainable_params)
-        obs = qml.dot(new_cotangents, new_observables)
-        if obs.pauli_rep is not None:
-            flat_bra = obs.pauli_rep.dot(ket.flatten(), wire_order=list(range(tape.num_wires)))
-            bra = 2 * flat_bra.reshape(ket.shape)
-        else:
-            bra = 2 * apply_operation(obs, ket)
 
         def real_if_expval(val):
             return np.real(val)
@@ -293,7 +380,13 @@ def adjoint_vjp(tape: QuantumTape, cotangents: Tuple[Number], state=None):
     param_number = len(tape.get_parameters(trainable_only=False, operations_only=True)) - 1
     trainable_param_number = len(tape.trainable_params) - 1
 
-    cotangents_in = np.empty(len(tape.trainable_params), dtype=tape.measurements[0].numeric_type)
+    res_shape = (
+        (len(tape.trainable_params),)
+        if batch_size is None
+        else (len(tape.trainable_params), batch_size)
+    )
+    cotangents_in = np.empty(res_shape, dtype=tape.measurements[0].numeric_type)
+    summing_axis = None if batch_size is None else tuple(range(1, np.ndim(bras)))
 
     for op in reversed(tape.operations[tape.num_preps :]):
         adj_op = qml.adjoint(op)
@@ -304,13 +397,15 @@ def adjoint_vjp(tape: QuantumTape, cotangents: Tuple[Number], state=None):
                 d_op_matrix = operation_derivative(op)
                 ket_temp = apply_operation(qml.QubitUnitary(d_op_matrix, wires=op.wires), ket)
 
-                cotangents_in[trainable_param_number] = real_if_expval(
-                    np.sum(np.conj(bra) * ket_temp)
-                )
+                # Pad cotangent in with zeros for batch number with zero cotangents
+                cot_in = real_if_expval(np.sum(np.conj(bras) * ket_temp, axis=summing_axis))
+                for i in null_batch_indices:
+                    cot_in = np.insert(cot_in, i, 0.0)
+                cotangents_in[trainable_param_number] = cot_in
 
                 trainable_param_number -= 1
             param_number -= 1
 
-        bra = apply_operation(adj_op, bra)
+        bras = apply_operation(adj_op, bras, is_state_batched=bool(batch_size))
 
     return tuple(cotangents_in)
diff --git a/pennylane/workflow/execution.py b/pennylane/workflow/execution.py
index efcf0fe1f..c8d8c37bf 100644
--- a/pennylane/workflow/execution.py
+++ b/pennylane/workflow/execution.py
@@ -137,7 +137,7 @@ def _get_ml_boundary_execute(
                 from .interfaces.jax_jit import jax_jit_jvp_execute as ml_boundary
         else:  # interface in {"jax", "jax-python", "JAX"}:
             if device_vjp:
-                from .interfaces.jax import jax_vjp_execute as ml_boundary
+                from .interfaces.jax_jit import jax_jit_vjp_execute as ml_boundary
             else:
                 from .interfaces.jax import jax_jvp_execute as ml_boundary
 
diff --git a/pennylane/workflow/interfaces/jax.py b/pennylane/workflow/interfaces/jax.py
index c302e3b14..0f9d9e04d 100644
--- a/pennylane/workflow/interfaces/jax.py
+++ b/pennylane/workflow/interfaces/jax.py
@@ -238,24 +238,9 @@ def _execute_and_compute_jvp(tapes, execute_fn, jpc, primals, tangents):
     return _to_jax(res), _to_jax(jvps)
 
 
-def _vjp_fwd(params, tapes, execute_fn, jpc):
-    """Perform the forward pass execution, return results and empty residuals."""
-    new_tapes = set_parameters_on_copy_and_unwrap(tapes.vals, params, unwrap=False)
-    return _to_jax(execute_fn(new_tapes)), None
-
-
-def _vjp_bwd(tapes, execute_fn, jpc, _, dy):
-    """Perform the backward pass of a vjp calculation, returning the vjp."""
-    vjp = jpc.compute_vjp(tapes.vals, dy)
-    return (_to_jax(vjp),)
-
-
 _execute_jvp = jax.custom_jvp(_execute_wrapper, nondiff_argnums=[1, 2, 3])
 _execute_jvp.defjvp(_execute_and_compute_jvp)
 
-_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, device=None):
     """Execute a batch of tapes with JAX parameters using JVP derivatives.
@@ -277,25 +262,3 @@ def jax_jvp_execute(tapes: Batch, execute_fn: ExecuteFn, jpc, device=None):
     parameters = tuple(tuple(t.get_parameters()) for t in tapes)
 
     return _execute_jvp(parameters, _NonPytreeWrapper(tuple(tapes)), execute_fn, 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:
-        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.
-
-    Returns:
-        TensorLike: A nested tuple of tape results. Each element in
-        the returned tuple corresponds in order to the provided tapes.
-
-    """
-    if logger.isEnabledFor(logging.DEBUG):  # pragma: no cover
-        logger.debug("Entry with (tapes=%s, execute_fn=%s, jpc=%s)", tapes, execute_fn, jpc)
-
-    parameters = tuple(tuple(t.get_parameters()) for t in tapes)
-
-    return _execute_vjp(parameters, _NonPytreeWrapper(tuple(tapes)), execute_fn, jpc)
diff --git a/pennylane/workflow/interfaces/jax_jit.py b/pennylane/workflow/interfaces/jax_jit.py
index eb87b84fe..1b41d6acc 100644
--- a/pennylane/workflow/interfaces/jax_jit.py
+++ b/pennylane/workflow/interfaces/jax_jit.py
@@ -65,12 +65,16 @@ def _to_jax(result: qml.typing.ResultBatch) -> qml.typing.ResultBatch:
     return jnp.array(result)
 
 
-def _set_parameters_on_copy(tapes, params):
-    """Copy a set of tapes with operations and set parameters"""
+def _set_all_parameters_on_copy(tapes, params):
+    """Copy a set of tapes with operations and set all parameters"""
     return tuple(t.bind_new_parameters(a, list(range(len(a)))) for t, a in zip(tapes, params))
 
 
-# pylint: disable=no-member
+def _set_trainable_parameters_on_copy(tapes, params):
+    """Copy a set of tapes with operations and set all trainable parameters"""
+    return tuple(t.bind_new_parameters(a, t.trainable_params) for t, a in zip(tapes, params))
+
+
 def _jax_dtype(m_type):
     if m_type == int:
         return jnp.int64 if jax.config.jax_enable_x64 else jnp.int32
@@ -138,7 +142,7 @@ def _execute_wrapper(params, tapes, execute_fn, _, device) -> ResultBatch:
     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)
+        new_tapes = _set_all_parameters_on_copy(tapes.vals, p)
         res = tuple(_to_jax(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
@@ -174,7 +178,7 @@ def _execute_and_compute_jvp(tapes, execute_fn, jpc, device, primals, tangents):
     )
 
     def wrapper(inner_params):
-        new_tapes = _set_parameters_on_copy(tapes.vals, inner_params)
+        new_tapes = _set_all_parameters_on_copy(tapes.vals, inner_params)
         return _to_jax(jpc.execute_and_compute_jacobian(new_tapes))
 
     res_struct = tuple(_result_shape_dtype_struct(t, device) for t in tapes.vals)
@@ -195,11 +199,11 @@ def _vjp_bwd(tapes, execute_fn, jpc, device, params, dy):
     """Perform the backward pass of a vjp calculation, returning the vjp."""
 
     def wrapper(inner_params, inner_dy):
-        new_tapes = _set_parameters_on_copy(tapes.vals, inner_params)
+        new_tapes = _set_trainable_parameters_on_copy(tapes.vals, inner_params)
         return _to_jax(jpc.compute_vjp(new_tapes, inner_dy))
 
     vjp_shape = _pytree_shape_dtype_struct(params)
-    return (jax.pure_callback(wrapper, vjp_shape, params, dy),)
+    return (jax.pure_callback(wrapper, vjp_shape, params, dy, vectorized=True),)
 
 
 _execute_jvp_jit = jax.custom_jvp(_execute_wrapper, nondiff_argnums=[1, 2, 3, 4])


Test patch:
diff --git a/tests/devices/qubit/test_adjoint_jacobian.py b/tests/devices/qubit/test_adjoint_jacobian.py
index c78949ebe..095074aff 100644
--- a/tests/devices/qubit/test_adjoint_jacobian.py
+++ b/tests/devices/qubit/test_adjoint_jacobian.py
@@ -470,7 +470,7 @@ class TestAdjointJVP:
 class TestAdjointVJP:
     """Test for adjoint_vjp"""
 
-    @pytest.mark.parametrize("cotangents", [(0,), (1.232,), 5.2])
+    @pytest.mark.parametrize("cotangents", [0, 1.232, 5.2])
     def test_single_param_single_obs(self, cotangents, tol):
         """Test VJP is correct for a single parameter and observable"""
         x = np.array(0.654)
@@ -478,8 +478,7 @@ class TestAdjointVJP:
 
         actual = adjoint_vjp(qs, cotangents)
 
-        iterable_cotangent = cotangents if isinstance(cotangents, tuple) else (cotangents,)
-        expected = -iterable_cotangent[0] * np.sin(x)
+        expected = -cotangents * np.sin(x)
         assert np.allclose(actual, expected, atol=tol)
 
     @pytest.mark.parametrize("cotangents", [(0, 0), (0, 0.653), (1.232, 2.963)])
@@ -497,7 +496,7 @@ class TestAdjointVJP:
         expected = np.dot(np.array([-np.sin(x), np.cos(x)]), np.array(cotangents))
         assert np.allclose(actual, expected, atol=tol)
 
-    @pytest.mark.parametrize("cotangents", [(0,), (1.232,)])
+    @pytest.mark.parametrize("cotangents", [0, 1.232])
     def test_multi_param_single_obs(self, cotangents, tol):
         """Test VJP is correct for multiple parameters and a single observable"""
         x = np.array(0.654)
@@ -511,7 +510,7 @@ class TestAdjointVJP:
         assert isinstance(actual, tuple)
         assert len(actual) == 2
 
-        expected = cotangents[0] * np.array([np.cos(x) * np.sin(y), np.sin(x) * np.cos(y)])
+        expected = cotangents * np.array([np.cos(x) * np.sin(y), np.sin(x) * np.cos(y)])
         assert np.allclose(actual, expected, atol=tol)
 
     @pytest.mark.parametrize(
@@ -564,6 +563,243 @@ class TestAdjointVJP:
         expected = np.array(cotangents) @ jac
         assert np.allclose(actual, expected, atol=tol)
 
+    @pytest.mark.parametrize(
+        "cotangents",
+        ((0, 1.23), (1.232, -2.098, 0.323, 1.112), (5.212, -0.354, -2.575), (0.0, 0.0, 0.0)),
+    )
+    def test_single_param_single_obs_batched(self, cotangents, tol):
+        """Test that batched cotangents with adjoint VJP give correct results when
+        the tape has a single trainable parameter and a single observable"""
+        x = np.array(0.654)
+        y = np.array(1.221)
+
+        obs = [qml.expval(qml.PauliZ(wires=[0]))]
+        qs = QuantumScript([qml.RY(x, wires=[0]), qml.RX(y, wires=[1])], obs, trainable_params=[0])
+
+        actual = adjoint_vjp(qs, cotangents)
+        assert isinstance(actual, tuple)
+        assert len(actual) == 1
+
+        jac = np.array([[-np.sin(x)]])
+        expected = jac.T @ np.expand_dims(np.array(cotangents), 0)
+        assert np.allclose(actual, expected, atol=tol)
+
+    @pytest.mark.parametrize(
+        "cotangents",
+        [
+            (np.array([1, 0, 0]), np.array([0, 1, 0]), np.array([0, 0, 1])),
+            (np.array([0.653, 0, 0]), np.array([0, 0.573, 0]), np.array([0, 0, 1.232])),
+            (np.array([0.653, -1.456]), np.array([0.498, 0.573]), np.array([0, 1.232])),
+            (
+                np.array([0.653, 0, 0, -1.234]),
+                np.array([-0.323, 0.573, -1.449, -0.573]),
+                np.array([0, 1, 1.232, 1.232]),
+            ),
+            (
+                np.array([0, 0, 0]),
+                np.array([0, 0, 0]),
+                np.array([0, 0, 0]),
+            ),
+        ],
+    )
+    def test_single_param_multi_obs_batched(self, cotangents, tol):
+        """Test that batched cotangents with adjoint VJP give correct results when
+        the tape has a single trainable parameter and multiple observables"""
+        x = np.array(0.654)
+        y = np.array(1.221)
+
+        obs = [
+            qml.expval(qml.PauliZ(wires=[0])),
+            qml.expval(qml.PauliY(wires=[1])),
+            qml.expval(qml.PauliX(wires=[0])),
+        ]
+        qs = QuantumScript([qml.RY(x, wires=[0]), qml.RX(y, wires=[1])], obs, trainable_params=[0])
+
+        actual = adjoint_vjp(qs, cotangents)
+        assert isinstance(actual, tuple)
+        assert len(actual) == 1
+
+        jac = np.array([[-np.sin(x)], [0], [np.cos(x)]])
+        expected = jac.T @ np.array(cotangents)
+        assert np.allclose(actual, expected, atol=tol)
+
+    @pytest.mark.parametrize(
+        "cotangents",
+        ((0, 1.23), (1.232, -2.098, 0.323, 1.112), (5.212, -0.354, -2.575), (0.0, 0.0, 0.0)),
+    )
+    def test_multi_param_single_obs_batched(self, cotangents, tol):
+        """Test that batched cotangents with adjoint VJP give correct results when
+        the tape has multiple trainable parameters and a single observable"""
+        x = np.array(0.654)
+        y = np.array(1.221)
+
+        obs = [qml.expval(qml.PauliZ(wires=[0]))]
+        qs = QuantumScript(
+            [qml.RY(x, wires=[0]), qml.RX(y, wires=[1])], obs, trainable_params=[0, 1]
+        )
+
+        actual = adjoint_vjp(qs, cotangents)
+        assert isinstance(actual, tuple)
+        assert len(actual) == 2
+
+        jac = np.array([[-np.sin(x), 0]])
+        expected = jac.T @ np.expand_dims(np.array(cotangents), 0)
+        assert np.allclose(actual, expected, atol=tol)
+
+    @pytest.mark.parametrize(
+        "cotangents",
+        [
+            (np.array([1, 0, 0]), np.array([0, 1, 0]), np.array([0, 0, 1])),
+            (np.array([0.653, 0, 0]), np.array([0, 0.573, 0]), np.array([0, 0, 1.232])),
+            (np.array([0.653, -1.456]), np.array([0.498, 0.573]), np.array([0, 1.232])),
+            (
+                np.array([0.653, 0, 0, -1.234]),
+                np.array([-0.323, 0.573, -1.449, -0.573]),
+                np.array([0, 1, 1.232, 1.232]),
+            ),
+            (
+                np.array([0, 0, 0]),
+                np.array([0, 0, 0]),
+                np.array([0, 0, 0]),
+            ),
+        ],
+    )
+    def test_multi_param_multi_obs_batched(self, cotangents, tol):
+        """Test that batched cotangents with adjoint VJP give correct results when
+        the tape has multiple trainable parameters and multiple observables"""
+        x = np.array(0.654)
+        y = np.array(1.221)
+
+        obs = [
+            qml.expval(qml.PauliZ(wires=[0])),
+            qml.expval(qml.PauliY(wires=[1])),
+            qml.expval(qml.PauliX(wires=[0])),
+        ]
+        qs = QuantumScript(
+            [qml.RY(x, wires=[0]), qml.RX(y, wires=[1])], obs, trainable_params=[0, 1]
+        )
+
+        actual = adjoint_vjp(qs, cotangents)
+        assert isinstance(actual, tuple)
+        assert len(actual) == 2
+
+        jac = np.array([[-np.sin(x), 0], [0, -np.cos(y)], [np.cos(x), 0]])
+        expected = jac.T @ np.array(cotangents)
+        assert np.allclose(actual, expected, atol=tol)
+
+    @pytest.mark.parametrize(
+        "cotangents",
+        [
+            (np.array([0.498, 0.573]), np.array([0.653, -1.456]), 0.0),
+            (np.array([0.498, 0.573, -1.456]), 0.0, 0.0),
+        ],
+    )
+    def test_inhomogenous_cotangents(self, cotangents, tol):
+        """Test that inhomogenous cotangents give the correct VJP"""
+        x = np.array(0.654)
+        y = np.array(1.221)
+
+        obs = [
+            qml.expval(qml.PauliZ(wires=[0])),
+            qml.expval(qml.PauliY(wires=[1])),
+            qml.expval(qml.PauliX(wires=[0])),
+        ]
+        qs = QuantumScript(
+            [qml.RY(x, wires=[0]), qml.RX(y, wires=[1])], obs, trainable_params=[0, 1]
+        )
+
+        actual = adjoint_vjp(qs, cotangents)
+        assert isinstance(actual, tuple)
+        assert len(actual) == 2
+
+        new_cotangents = []
+        inner_shape = len(cotangents[0])
+        for c in cotangents:
+            if isinstance(c, float):
+                new_cotangents.append(np.zeros(inner_shape))
+            else:
+                new_cotangents.append(c)
+
+        jac = np.array([[-np.sin(x), 0], [0, -np.cos(y)], [np.cos(x), 0]])
+        expected = jac.T @ np.array(new_cotangents)
+        assert np.allclose(actual, expected, atol=tol)
+
+    @pytest.mark.parametrize(
+        "cotangents",
+        [
+            np.array(
+                [
+                    [0.123, 0.456, 0.789, -0.123],
+                    [-0.456, -0.789, 1.234, 5.678],
+                    [-0.345, -4.345, -2.589, 3.456],
+                ]
+            ),
+            np.array([[0.0, 0.123, 0.765, 4.123], [-7.698, -3.465, -1.289, 4.697]]),
+        ],
+    )
+    def test_single_param_state_batched(self, cotangents, tol):
+        """Test that computing the VJP with batched cotangents for state measurements
+        gives the correct results for a single trainable parameter."""
+        x = np.array(0.654)
+        y = np.array(1.221)
+
+        obs = [qml.state()]
+        qs = QuantumScript([qml.RY(x, wires=[0]), qml.RX(y, wires=[1])], obs, trainable_params=[0])
+
+        actual = adjoint_vjp(qs, cotangents)
+        assert isinstance(actual, tuple)
+        assert len(actual) == 1
+
+        jac = np.array(
+            [
+                [-0.5 * np.sin(x / 2) * np.cos(y / 2)],
+                [0.5j * np.sin(x / 2) * np.sin(y / 2)],
+                [0.5 * np.cos(x / 2) * np.cos(y / 2)],
+                [-0.5j * np.cos(x / 2) * np.sin(y / 2)],
+            ]
+        )
+        expected = jac.T @ cotangents.T
+        assert np.allclose(actual, expected, atol=tol)
+
+    @pytest.mark.parametrize(
+        "cotangents",
+        [
+            np.array(
+                [
+                    [0.123, 0.456, 0.789, -0.123],
+                    [-0.456, -0.789, 1.234, 5.678],
+                    [-0.345, -4.345, -2.589, 3.456],
+                ]
+            ),
+            np.array([[0.0, 0.123, 0.765, 4.123], [-7.698, -3.465, -1.289, 4.697]]),
+        ],
+    )
+    def test_multi_param_state_batched(self, cotangents, tol):
+        """Test that computing the VJP with batched cotangents for state measurements
+        gives the correct results for multiple trainable parameters."""
+        x = np.array(0.654)
+        y = np.array(1.221)
+
+        obs = [qml.state()]
+        qs = QuantumScript(
+            [qml.RY(x, wires=[0]), qml.RX(y, wires=[1])], obs, trainable_params=[0, 1]
+        )
+
+        actual = adjoint_vjp(qs, cotangents)
+        assert isinstance(actual, tuple)
+        assert len(actual) == 2
+
+        jac = np.array(
+            [
+                [-0.5 * np.sin(x / 2) * np.cos(y / 2), -0.5 * np.cos(x / 2) * np.sin(y / 2)],
+                [0.5j * np.sin(x / 2) * np.sin(y / 2), -0.5j * np.cos(x / 2) * np.cos(y / 2)],
+                [0.5 * np.cos(x / 2) * np.cos(y / 2), -0.5 * np.sin(x / 2) * np.sin(y / 2)],
+                [-0.5j * np.cos(x / 2) * np.sin(y / 2), -0.5j * np.sin(x / 2) * np.cos(y / 2)],
+            ]
+        )
+        expected = jac.T @ cotangents.T
+        assert np.allclose(actual, expected, atol=tol)
+
     def test_with_nontrainable_parametrized(self):
         """Test that a parametrized `QubitUnitary` is accounted for correctly
         when it is not trainable."""
diff --git a/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py
index d8f2c0f0e..770d8112d 100644
--- a/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py
@@ -202,10 +202,6 @@ class TestJaxExecuteIntegration:
 
     def test_jacobian(self, execute_kwargs, shots, device):
         """Test jacobian calculation"""
-        if execute_kwargs.get("gradient_fn", "") == "adjoint" and execute_kwargs.get(
-            "device_vjp", False
-        ):
-            pytest.xfail("adjoint_vjp does not yet support jax jacobians.")
 
         a = jnp.array(0.1)
         b = jnp.array(0.2)
@@ -288,11 +284,6 @@ class TestJaxExecuteIntegration:
     def test_tapes_with_different_return_size(self, execute_kwargs, shots, device):
         """Test that tapes wit different can be executed and differentiated."""
 
-        if execute_kwargs.get("gradient_fn", "") == "adjoint" and execute_kwargs.get(
-            "device_vjp", False
-        ):
-            pytest.xfail("adjoint_vjp does not yet support jax jacobians.")
-
         def cost(params):
             tape1 = qml.tape.QuantumScript(
                 [qml.RY(params[0], 0), qml.RX(params[1], 0)],
@@ -387,14 +378,6 @@ class TestJaxExecuteIntegration:
             return jnp.hstack(execute([new_tape], device, **execute_kwargs)[0])
 
         jac_fn = jax.jacobian(cost, argnums=[0, 1])
-        if execute_kwargs.get("gradient_fn", "") == "adjoint" and execute_kwargs.get(
-            "device_vjp", False
-        ):
-            with pytest.raises(
-                NotImplementedError, match=r"adjoint_vjp does not yet support JAX Jacobians"
-            ):
-                jac = jac_fn(a, b)
-            return
         jac = jac_fn(a, b)
 
         a = jnp.array(0.54)
@@ -564,10 +547,6 @@ class TestJaxExecuteIntegration:
             assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0)
 
         jac_fn = jax.jacobian(cost, argnums=[0, 1])
-        if execute_kwargs.get("device_vjp", False):
-            with pytest.raises(NotImplementedError):
-                jac_fn(x, y)
-            return
         res = jac_fn(x, y)
         assert isinstance(res, tuple) and len(res) == 2
         assert res[0].shape == (2, 4) if shots.has_partitioned_shots else (4,)
@@ -625,10 +604,6 @@ class TestJaxExecuteIntegration:
             assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0)
 
         jac_fn = jax.jacobian(cost, argnums=[0, 1])
-        if execute_kwargs.get("device_vjp", False):
-            with pytest.raises(NotImplementedError):
-                jac_fn(x, y)
-            return
         res = jac_fn(x, y)
         assert isinstance(res, tuple) and len(res) == 2
         assert res[0].shape == (3 * shots.num_copies,) if shots.has_partitioned_shots else (3,)
@@ -796,8 +771,6 @@ class TestHamiltonianWorkflows:
 
         if execute_kwargs["gradient_fn"] == "adjoint" and not use_new_op_math:
             pytest.skip("adjoint differentiation does not suppport hamiltonians.")
-        if execute_kwargs["gradient_fn"] == "adjoint" and execute_kwargs.get("device_vjp", False):
-            pytest.xfail("adjoint vjp does not yet support jax jacobians.")
 
         coeffs1 = jnp.array([0.1, 0.2, 0.3])
         coeffs2 = jnp.array([0.7])
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 33e6239f9..f8ded4f49 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
@@ -424,9 +424,7 @@ class TestVectorValuedQNode:
         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")
-        elif diff_method == "spsa":
+        if diff_method == "spsa":
             gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
             gradient_kwargs["num_directions"] = 20
             tol = TOL_FOR_SPSA
@@ -477,9 +475,7 @@ class TestVectorValuedQNode:
         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")
-        elif diff_method == "spsa":
+        if diff_method == "spsa":
             gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
             gradient_kwargs["num_directions"] = 20
             tol = TOL_FOR_SPSA
@@ -563,9 +559,7 @@ class TestVectorValuedQNode:
         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")
-        elif diff_method == "spsa":
+        if diff_method == "spsa":
             gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
             gradient_kwargs["num_directions"] = 20
             tol = TOL_FOR_SPSA
@@ -639,9 +633,7 @@ class TestVectorValuedQNode:
         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")
-        elif diff_method == "spsa":
+        if diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
             tol = TOL_FOR_SPSA
 
@@ -692,9 +684,7 @@ class TestVectorValuedQNode:
         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")
-        elif diff_method == "hadamard":
+        if diff_method == "hadamard":
             pytest.skip("Hadamard does not support var")
         elif diff_method == "spsa":
             gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
@@ -1355,8 +1345,6 @@ class TestQubitIntegrationHigherOrder:
 
     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")
 
         x = jax.numpy.array(0.543)
         y = jax.numpy.array(-0.654)
@@ -1872,10 +1860,10 @@ class TestJIT:
         """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 does not support probs")
         if dev.name == "param_shift.qubit":
             pytest.xfail("gradient transforms have a different vjp shape convention.")
+        elif jacobian == jax.jacfwd and device_vjp:
+            pytest.skip("device vjps are not compatible with forward differentiation.")
         elif diff_method == "spsa":
             gradient_kwargs = {"h": H_FOR_SPSA, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)}
             tol = TOL_FOR_SPSA
@@ -2065,9 +2053,8 @@ class TestReturn:
             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 diff_method == "adjoint":
-            pytest.skip("Test does not supports adjoint because of probabilities.")
+        if jacobian == jax.jacfwd and device_vjp:
+            pytest.skip("jacfwd is not compatible with device_vjp=True.")
 
         @qnode(
             dev,
@@ -2095,10 +2082,10 @@ class TestReturn:
         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")
+        if jacobian == jax.jacfwd and device_vjp:
+            pytest.skip("jacfwd is not compatible with device_vjp=True.")
 
         @qnode(
             dev,
@@ -2132,10 +2119,10 @@ class TestReturn:
         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")
+        if jacobian == jax.jacfwd and device_vjp:
+            pytest.skip("jacfwd is not compatible with device_vjp=True.")
 
         @qnode(
             dev,
@@ -2244,12 +2231,14 @@ class TestReturn:
         """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":
             pytest.skip("Test does not supports hadamard because of var.")
         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.")
+        if diff_method == "adjoint":
+            pytest.skip("adjoint supports either all expvals or only diagonal measurements")
 
         par_0 = jax.numpy.array(0.1)
         par_1 = jax.numpy.array(0.2)
@@ -2294,12 +2283,14 @@ class TestReturn:
         """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":
             pytest.skip("Test does not supports hadamard because of var.")
         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.")
+        if diff_method == "adjoint":
+            pytest.skip("adjoint supports either all expvals or only diagonal measurements")
 
         @qnode(
             dev,
@@ -2337,9 +2328,6 @@ class TestReturn:
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
 
-        if diff_method == "adjoint":
-            pytest.skip("Test does not supports adjoint because of probabilities.")
-
         @qnode(
             dev,
             interface=interface,
@@ -2371,10 +2359,10 @@ class TestReturn:
         """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")
+        if jacobian == jax.jacfwd and device_vjp:
+            pytest.skip("jacfwd is not compatible with device_vjp=True.")
 
         @qnode(
             dev,
@@ -2416,10 +2404,10 @@ class TestReturn:
         """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")
+        if jacobian == jax.jacfwd and device_vjp:
+            pytest.skip("jacfwd is not compatible with device_vjp=True.")
 
         @qnode(
             dev,
diff --git a/tests/interfaces/default_qubit_2_integration/test_jax_qnode_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_jax_qnode_default_qubit_2.py
index 1f74286ca..20b28b7ee 100644
--- a/tests/interfaces/default_qubit_2_integration/test_jax_qnode_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_jax_qnode_default_qubit_2.py
@@ -296,10 +296,6 @@ class TestVectorValuedQNode:
         assert np.allclose(res[0], expected[0], atol=tol, rtol=0)
         assert np.allclose(res[1], expected[1], atol=tol, rtol=0)
 
-        if device_vjp:
-            with pytest.raises(NotImplementedError):
-                jax.jacobian(circuit, argnums=[0, 1])(a, b)
-            return
         res = jax.jacobian(circuit, argnums=[0, 1])(a, b)
         expected = np.array([[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]])
         assert circuit.qtape.trainable_params == [0, 1]
@@ -347,10 +343,6 @@ class TestVectorValuedQNode:
             qml.CNOT(wires=[0, 1])
             return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))
 
-        if device_vjp:
-            with pytest.raises(NotImplementedError):
-                jax.jacobian(circuit, argnums=[0, 1])(a, b)
-            return
         jac_fn = jax.jacobian(circuit, argnums=[0, 1])
         res = jac_fn(a, b)
 
@@ -404,10 +396,6 @@ class TestVectorValuedQNode:
             qml.CNOT(wires=[0, 1])
             return qml.probs(wires=[1])
 
-        if device_vjp:
-            with pytest.raises(NotImplementedError):
-                jax.jacobian(circuit, argnums=[0, 1])(x, y)
-            return
         res = jax.jacobian(circuit, argnums=[0, 1])(x, y)
 
         expected = np.array(
@@ -473,10 +461,6 @@ class TestVectorValuedQNode:
         assert res[1].shape == (4,)  # pylint:disable=comparison-with-callable
         assert np.allclose(res[1], expected[1], atol=tol, rtol=0)
 
-        if device_vjp:
-            with pytest.raises(NotImplementedError):
-                jax.jacobian(circuit, argnums=[0, 1])(x, y)
-            return
         jac = jax.jacobian(circuit, argnums=[0, 1])(x, y)
         expected_0 = np.array(
             [
@@ -551,10 +535,6 @@ class TestVectorValuedQNode:
         assert res[1].shape == (2,)  # pylint:disable=comparison-with-callable
         assert np.allclose(res[1], expected[1], atol=tol, rtol=0)
 
-        if device_vjp:
-            with pytest.raises(NotImplementedError):
-                jax.jacobian(circuit, argnums=[0, 1])(x, y)
-            return
         jac = jax.jacobian(circuit, argnums=[0, 1])(x, y)
         expected = [
             [-np.sin(x), 0],
@@ -616,10 +596,6 @@ class TestVectorValuedQNode:
             qml.CNOT(wires=[0, 1])
             return qml.expval(qml.PauliZ(0)), qml.probs(wires=[1])
 
-        if device_vjp:
-            with pytest.raises(NotImplementedError):
-                jax.jacobian(circuit, argnums=[0])(x, y)
-            return
         jac = jax.jacobian(circuit, argnums=[0])(x, y)
 
         expected = [
@@ -685,10 +661,6 @@ class TestVectorValuedQNode:
         assert res[1].shape == (2,)  # pylint:disable=comparison-with-callable
         assert np.allclose(res[1], expected[1], atol=tol, rtol=0)
 
-        if device_vjp:
-            with pytest.raises(NotImplementedError):
-                jax.jacobian(circuit, argnums=[0, 1])(x, y)
-            return
         jac = jax.jacobian(circuit, argnums=[0, 1])(x, y)
         expected = [
             [2 * np.cos(x) * np.sin(x), 0],
@@ -1739,10 +1711,6 @@ class TestReturn:  # pylint:disable=too-many-public-methods
             qml.CNOT(wires=[0, 1])
             return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0))
 
-        if device_vjp:
-            with pytest.raises(NotImplementedError):
-                jacobian(circuit, argnums=[0, 1])(par_0, par_1, shots=shots)
-            return
         jac = jacobian(circuit, argnums=[0, 1])(par_0, par_1, shots=shots)
 
         assert isinstance(jac, tuple)
@@ -1785,10 +1753,6 @@ class TestReturn:  # pylint:disable=too-many-public-methods
 
         a = jax.numpy.array([0.1, 0.2])
 
-        if device_vjp:
-            with pytest.raises(NotImplementedError):
-                jacobian(circuit)(a, shots=shots)
-            return
         jac = jacobian(circuit)(a, shots=shots)
 
         assert isinstance(jac, tuple)
@@ -1910,10 +1874,6 @@ class TestReturn:  # pylint:disable=too-many-public-methods
 
         a = jax.numpy.array(0.1)
 
-        if device_vjp:
-            with pytest.raises(NotImplementedError):
-                jacobian(circuit)(a, shots=shots)
-            return
         jac = jacobian(circuit)(a, shots=shots)
 
         assert isinstance(jac, tuple)
@@ -1951,10 +1911,6 @@ class TestReturn:  # pylint:disable=too-many-public-methods
         a = jax.numpy.array(0.1)
         b = jax.numpy.array(0.2)
 
-        if device_vjp:
-            with pytest.raises(NotImplementedError):
-                jacobian(circuit, argnums=[0, 1])(a, b, shots=shots)
-            return
         jac = jacobian(circuit, argnums=[0, 1])(a, b, shots=shots)
 
         assert isinstance(jac, tuple)
@@ -1999,10 +1955,6 @@ class TestReturn:  # pylint:disable=too-many-public-methods
 
         a = jax.numpy.array([0.1, 0.2])
 
-        if device_vjp:
-            with pytest.raises(NotImplementedError):
-                jacobian(circuit)(a, shots=shots)
-            return
         jac = jacobian(circuit)(a, shots=shots)
 
         assert isinstance(jac, tuple)
diff --git a/tests/interfaces/test_jax_jit_qnode.py b/tests/interfaces/test_jax_jit_qnode.py
index 63a6000e7..7b9a3cc4c 100644
--- a/tests/interfaces/test_jax_jit_qnode.py
+++ b/tests/interfaces/test_jax_jit_qnode.py
@@ -1775,9 +1775,7 @@ class TestJIT:
         dev = qml.device(dev_name, wires=num_wires)
 
         gradient_kwargs = {}
-        if diff_method == "adjoint":
-            pytest.xfail(reason="The adjoint method is not using host-callback currently")
-        elif diff_method == "spsa":
+        if diff_method == "spsa":
             gradient_kwargs = {"h": H_FOR_SPSA, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)}
             tol = TOL_FOR_SPSA
 
@@ -1919,7 +1917,7 @@ class TestJIT:
 
         gradient_kwargs = {}
         if diff_method == "adjoint":
-            pytest.xfail(reason="The adjoint method is not using host-callback currently")
+            pytest.xfail(reason="Adjoint does not support probs.")
         elif diff_method == "spsa":
             gradient_kwargs = {"h": H_FOR_SPSA, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)}
             tol = TOL_FOR_SPSA