Repository: PennyLaneAI/pennylane
Pull request #4865 (2023-11-21 17:19:20)

Add state measurement support for adjoint jac and adjoint vjp
This draft adds support to `default.qubit`'s adjoint jacobian and adjoint vjp for state measurements.


```
dev = qml.device('default.qubit')

@qml.qnode(dev, diff_method="adjoint", device_vjp=True)
def circuit(x):
    qml.IsingXX(x, wires=(0,1))
    return qml.vn_entropy(wires=0)

qml.grad(circuit)(qml.numpy.array(0.2))
```
```
0.45678925885706456
```

[sc-50755]

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.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4865?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 [(`a05ae28`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/a05ae286c5925e0c881099cea751193d368779fc?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.66% compared to head [(`aefde35`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4865?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    #4865      +/-   ##
==========================================
- Coverage   99.66%   99.66%   -0.01%     
==========================================
  Files         388      388              
  Lines       35362    35154     -208     
==========================================
- Hits        35245    35036     -209     
- Misses        117      118       +1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4865?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).


Patch:
diff --git a/doc/introduction/interfaces.rst b/doc/introduction/interfaces.rst
index 8b97bdbd9..487da2e7a 100644
--- a/doc/introduction/interfaces.rst
+++ b/doc/introduction/interfaces.rst
@@ -418,7 +418,7 @@ At the moment, it takes into account the following parameters:
 +                  +------------------------------+--------------+---------------+--------------+--------------+---------------+----------------+----------------+-------------+-------------+-------------+
 |                  | ``"backprop"``               |     :gr:`4`  |   :gr:`4`     |     :gr:`5`  |     :rd:`9`  |   :gr:`5`     |    :gr:`5`     |   :gr:`5`      | :gr:`5`     | :gr:`5`     | :gr:`5`     |
 +                  +------------------------------+--------------+---------------+--------------+--------------+---------------+----------------+----------------+-------------+-------------+-------------+
-|                  | ``"adjoint"``                |      :rd:`6` |     :rd:`6`   |  :rd:`6`     | :rd:`6`      |      :gr:`7`  |  :gr:`7`       |   :gr:`7`      | :rd:`6`     | :rd:`6`     | :rd:`6`     |
+|                  | ``"adjoint"``                |      :gr:`7` |     :gr:`7`   |  :gr:`7`     | :rd:`6`      |      :gr:`7`  |  :gr:`7`       |   :gr:`7`      | :gr:`7`     | :gr:`7`     | :gr:`7`     |
 +                  +------------------------------+--------------+---------------+--------------+--------------+---------------+----------------+----------------+-------------+-------------+-------------+
 |                  | ``"parameter-shift"``        |   :rd:`10`   |    :rd:`10`   |   :gr:`8`    |  :rd:`9`     |   :gr:`8`     |   :gr:`8`      | :gr:`8`        |   :gr:`8`   |   :rd:`10`  |   :rd:`10`  |
 +                  +------------------------------+--------------+---------------+--------------+--------------+---------------+----------------+----------------+-------------+-------------+-------------+
@@ -432,7 +432,7 @@ At the moment, it takes into account the following parameters:
 +                  +------------------------------+--------------+---------------+--------------+--------------+---------------+----------------+----------------+-------------+-------------+-------------+
 |                  | ``"backprop"``               |     :gr:`5`  |   :gr:`5`     |     :gr:`5`  |     :rd:`9`  |   :gr:`5`     |    :gr:`5`     |   :gr:`5`      | :gr:`5`     | :gr:`5`     | :gr:`5`     |
 +                  +------------------------------+--------------+---------------+--------------+--------------+---------------+----------------+----------------+-------------+-------------+-------------+
-|                  | ``"adjoint"``                |      :rd:`6` |     :rd:`6`   |  :rd:`6`     | :rd:`6`      |      :gr:`7`  |  :gr:`7`       |   :gr:`7`      | :rd:`6`     | :rd:`6`     | :rd:`6`     |
+|                  | ``"adjoint"``                |      :gr:`7` |     :gr:`7`   |  :gr:`7`     | :rd:`6`      |      :gr:`7`  |  :gr:`7`       |   :gr:`7`      | :gr:`7`     | :gr:`7`     | :gr:`7`     |
 +                  +------------------------------+--------------+---------------+--------------+--------------+---------------+----------------+----------------+-------------+-------------+-------------+
 |                  | ``"parameter-shift"``        |   :rd:`10`   |    :rd:`10`   |   :gr:`8`    |  :rd:`9`     |   :gr:`8`     |   :gr:`8`      | :gr:`8`        |   :gr:`8`   |   :rd:`10`  |   :rd:`10`  |
 +                  +------------------------------+--------------+---------------+--------------+--------------+---------------+----------------+----------------+-------------+-------------+-------------+
@@ -446,7 +446,7 @@ At the moment, it takes into account the following parameters:
 +                  +------------------------------+--------------+---------------+--------------+--------------+---------------+----------------+----------------+-------------+-------------+-------------+
 |                  | ``"backprop"``               |     :gr:`5`  |   :gr:`5`     |     :gr:`5`  |     :rd:`9`  |   :gr:`5`     |    :gr:`5`     |   :gr:`5`      | :gr:`5`     | :gr:`5`     | :gr:`5`     |
 +                  +------------------------------+--------------+---------------+--------------+--------------+---------------+----------------+----------------+-------------+-------------+-------------+
-|                  | ``"adjoint"``                |      :rd:`6` |     :rd:`6`   |  :rd:`6`     | :rd:`6`      |      :gr:`7`  |  :gr:`7`       |   :gr:`7`      | :rd:`6`     | :rd:`6`     | :rd:`6`     |
+|                  | ``"adjoint"``                |      :gr:`7` |     :gr:`7`   |  :rd:`6`     | :rd:`6`      |      :gr:`7`  |  :gr:`7`       |   :gr:`7`      | :gr:`7`     | :gr:`7`     | :gr:`7`     |
 +                  +------------------------------+--------------+---------------+--------------+--------------+---------------+----------------+----------------+-------------+-------------+-------------+
 |                  | ``"parameter-shift"``        |   :rd:`10`   |    :rd:`10`   |   :gr:`8`    |  :rd:`9`     |   :gr:`8`     |   :gr:`8`      | :gr:`8`        |   :gr:`8`   |   :rd:`10`  |   :rd:`10`  |
 +                  +------------------------------+--------------+---------------+--------------+--------------+---------------+----------------+----------------+-------------+-------------+-------------+
@@ -460,7 +460,7 @@ At the moment, it takes into account the following parameters:
 +                  +------------------------------+--------------+---------------+--------------+--------------+---------------+----------------+----------------+-------------+-------------+-------------+
 |                  | ``"backprop"``               |     :gr:`5`  |   :gr:`5`     |     :gr:`5`  |     :rd:`9`  |   :gr:`5`     |    :gr:`5`     |   :gr:`5`      | :gr:`5`     | :gr:`5`     | :gr:`5`     |
 +                  +------------------------------+--------------+---------------+--------------+--------------+---------------+----------------+----------------+-------------+-------------+-------------+
-|                  | ``"adjoint"``                |      :rd:`6` |     :rd:`6`   |  :rd:`6`     | :rd:`6`      |      :gr:`7`  |  :gr:`7`       |   :gr:`7`      | :rd:`6`     | :rd:`6`     | :rd:`6`     |
+|                  | ``"adjoint"``                |      :gr:`7` |     :gr:`7`   |  :gr:`7`     | :rd:`6`      |      :gr:`7`  |  :gr:`7`       |   :gr:`7`      | :gr:`7`     | :gr:`7`     | :gr:`7`     |
 +                  +------------------------------+--------------+---------------+--------------+--------------+---------------+----------------+----------------+-------------+-------------+-------------+
 |                  | ``"parameter-shift"``        |   :rd:`10`   |    :rd:`10`   |   :gr:`8`    |  :rd:`9`     |   :gr:`8`     |   :gr:`8`      | :gr:`8`        |   :gr:`8`   |   :rd:`10`  |   :rd:`10`  |
 +                  +------------------------------+--------------+---------------+--------------+--------------+---------------+----------------+----------------+-------------+-------------+-------------+
@@ -482,9 +482,9 @@ At the moment, it takes into account the following parameters:
    directly. However, any real scalar-valued post-processing done to the output of the
    circuit will be differentiable. See :ref:`State gradients <State gradients>` for details.
 5. Supported, but only when ``shots=None``. See :ref:`Backpropagation <Analytic backpropagation>` for details.
-6. Not supported. The adjoint differentiation algorithm is only implemented for computing the expectation values of observables. See
+6. Not supported. The adjoint differentiation algorithm is only implemented for analytic simulation. See
    :ref:`Adjoint differentation <Adjoint differentation>` for details.
-7. Supported. Raises warning when ``shots>0`` since the gradient is always computed analytically. See
+7. Supported. Raises error when ``shots>0`` since the gradient is always computed analytically. See
    :ref:`Adjoint differentation <Adjoint differentation>` for details.
 8. Supported.
 9. Not supported. The discretization of the output caused by wave function collapse is
diff --git a/doc/introduction/unsupported_gradients.rst b/doc/introduction/unsupported_gradients.rst
index f4d75ffe6..6a83e7888 100644
--- a/doc/introduction/unsupported_gradients.rst
+++ b/doc/introduction/unsupported_gradients.rst
@@ -96,8 +96,7 @@ Adjoint differentiation
 
 PennyLane implements the adjoint differentiation method from
 `2009.02823 <https://arxiv.org/pdf/2009.02823.pdf>`__, which only discusses
-the gradient of expectation of observables. The implementation is specific to the paper, hence the return statement
-of the quantum function wrapped in ``qml.qnode`` can only contain :func:`~.pennylane.expval` as a measurement.
+the gradient of expectation values of observables.
 
 In particular, the following code works as expected:
 
@@ -117,34 +116,30 @@ In particular, the following code works as expected:
 >>> print_grad()
 [-0.09983342]
 
-But the following code raises an error:
+``default.qubit`` can differentiate any other measurement process as long as it
+is in the Z measurement basis. In this case, we recommend using the device-provided vjp
+(``device_vjp=True``) for improved performance scaling. This algorithm works
+best when the final cost function only has a scalar value.
 
-.. code-block:: python
-
-    def print_grad_bad():
-        dev = qml.device('default.qubit', wires=1, shots=None)
+``lightning.qubit`` only supports expectation values.
 
-        @qml.qnode(dev, diff_method='adjoint')
-        def circuit(x):
-            qml.RX(x[0], wires=0)
-            return qml.state()
+.. code-block:: python 
 
-        def cost_fn(x):
-            out = circuit(x)
-            return np.abs(out[0])
+    @qml.qnode(qml.device('default.qubit'), diff_method="adjoint", device_vjp=True)
+    def circuit(x):
+        qml.IsingXX(x, wires=(0,1))
+        return qml.probs(wires=(0,1))
 
-        x = np.array([0.1], requires_grad=True)
-        print(qml.grad(cost_fn)(x))
+    def cost(x):
+        probs = circuit(x)
+        target = np.array([0, 0, 0, 1])
+        return qml.math.norm(probs-target)
 
->>> print_grad_bad()
-Traceback (most recent call last):
-  ...
-  File "C:\pennylane\pennylane\_qubit_device.py", line 951, in adjoint_jacobian
-    raise qml.QuantumFunctionError(
-pennylane.QuantumFunctionError: Adjoint differentiation method does not support measurement state
+>>> qml.grad(cost)(qml.numpy.array(0.1))
+-0.07059288589999416
 
-Furthermore, the adjoint differentiation algorithm is analytic by nature. If the user creates a device
-with ``shots>0``, a warning is raised and gradients are computed analytically:
+Furthermore, the adjoint differentiation algorithm is analytic by nature. If the an execution
+has ``shots>0``, an error is raised:
 
 .. code-block:: python
 
@@ -160,11 +155,7 @@ with ``shots>0``, a warning is raised and gradients are computed analytically:
         print(qml.grad(circuit)(x))
 
 >>> print_grad_ok()
-C:\pennylane\pennylane\qnode.py:434: UserWarning: Requested adjoint differentiation to be computed with finite shots. Adjoint differentiation always calculated exactly.
-  warnings.warn(
-C:\pennylane\pennylane\_qubit_device.py:965: UserWarning: Requested adjoint differentiation to be computed with finite shots. The derivative is always exact when using the adjoint differentiation method.
-  warnings.warn(
-[-0.09983342]
+DeviceError: Finite shots are not supported with adjoint + default.qubit
 
 .. _State gradients:
 
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index a4cfb3f3b..26a8ee841 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -384,6 +384,9 @@
 
 <h4>Other improvements</h4>
 
+* `default.qubit` now supports adjoint differentiation for arbitrary diagonal state-based measurements.
+  [(#4865)](https://github.com/PennyLaneAI/pennylane/pull/4865)
+
 * `qml.quantum_monte_carlo` now uses the new transform system.
   [(#4708)](https://github.com/PennyLaneAI/pennylane/pull/4708/)
 
diff --git a/pennylane/devices/default_qubit.py b/pennylane/devices/default_qubit.py
index cf000d435..1e0baaec0 100644
--- a/pennylane/devices/default_qubit.py
+++ b/pennylane/devices/default_qubit.py
@@ -104,6 +104,48 @@ def accepted_sample_measurement(m: qml.measurements.MeasurementProcess) -> bool:
     )
 
 
+def null_postprocessing(results):
+    """An empty post-processing function."""
+    return results[0]
+
+
+def all_state_postprocessing(results, measurements, wire_order):
+    """Process a state measurement back into the original measurements."""
+    result = tuple(m.process_state(results[0], wire_order=wire_order) for m in measurements)
+    return result[0] if len(measurements) == 1 else result
+
+
+@qml.transform
+def adjoint_state_measurements(tape: QuantumTape) -> (Tuple[QuantumTape], Callable):
+    """Perform adjoint measurement preprocessing.
+
+    * Allows a tape with only expectation values through unmodified
+    * Raises an error if non-expectation value measurements exist and any have diagonalizing gates
+    * Turns the circuit into a state measurement + classical postprocesssing for arbitrary measurements
+
+    Args:
+        tape (QuantumTape): the input circuit
+
+    """
+    if all(isinstance(m, qml.measurements.ExpectationMP) for m in tape.measurements):
+        return (tape,), null_postprocessing
+
+    if any(len(m.diagonalizing_gates()) > 0 for m in tape.measurements):
+        raise qml.DeviceError(
+            "adjoint diff supports either all expectation values or only measurements without observables."
+        )
+
+    params = tape.get_parameters()
+    complex_data = [qml.math.cast(p, complex) for p in params]
+    tape = tape.bind_new_parameters(complex_data, list(range(len(params))))
+    state_tape = qml.tape.QuantumScript(
+        tape.operations, [qml.measurements.StateMP(wires=tape.wires)]
+    )
+    return (state_tape,), partial(
+        all_state_postprocessing, measurements=tape.measurements, wire_order=tape.wires
+    )
+
+
 def _add_adjoint_transforms(program: TransformProgram) -> None:
     """Private helper function for ``preprocess`` that adds the transforms specific
     for adjoint differentiation.
@@ -124,9 +166,6 @@ def _add_adjoint_transforms(program: TransformProgram) -> None:
         """Specifies whether or not an observable is compatible with adjoint differentiation on DefaultQubit."""
         return obs.has_matrix
 
-    def accepted_adjoint_measurement(m: qml.measurements.MeasurementProcess) -> bool:
-        return isinstance(m, qml.measurements.ExpectationMP)
-
     name = "adjoint + default.qubit"
     program.add_transform(no_sampling, name=name)
     program.add_transform(
@@ -137,9 +176,9 @@ def _add_adjoint_transforms(program: TransformProgram) -> None:
     program.add_transform(validate_observables, adjoint_observables, name=name)
     program.add_transform(
         validate_measurements,
-        analytic_measurements=accepted_adjoint_measurement,
         name=name,
     )
+    program.add_transform(adjoint_state_measurements)
     program.add_transform(qml.transforms.broadcast_expand)
     program.add_transform(warn_about_trainable_observables)
 
diff --git a/pennylane/devices/qubit/adjoint_jacobian.md b/pennylane/devices/qubit/adjoint_jacobian.md
new file mode 100644
index 000000000..4d6c2dd77
--- /dev/null
+++ b/pennylane/devices/qubit/adjoint_jacobian.md
@@ -0,0 +1,158 @@
+# The adjoint jacobian algorithm
+
+## Full jacobian with expectation values
+
+For a statevector,
+
+$$
+| \Psi \rangle = U_n(x_n) \dots U_i (x_i) \dots U_0(x_0) |0\rangle
+$$
+
+we have the derivative of an expectation value of an observable $M$,
+
+$$
+\frac{\partial}{\partial x_i} \langle \Psi | M | \Psi \rangle
+$$
+
+$$
+= \langle 0| \dots \frac{\partial U_i^{\dagger}}{\partial x_i} \dots M \dots
+U_i \dots |0\rangle +
+\langle 0 | \dots U_i^{\dagger} \dots M \dots \frac{\partial U_i}{\partial x_i} \dots |0 \rangle
+$$
+
+$$
+= 2 \text{Re} \big( \langle \Psi | M U_n \dots \frac{\partial U_i}{\partial x_i}  \dots U_0 |0\rangle \big)
+$$
+
+We can then make the definitions:
+
+$$
+\langle b_i|  = 2 \langle \Psi | M U_n \dots U_{i+1}
+$$
+
+$$
+|k_i \rangle  = U_{i-1} \dots U_0 |0\rangle
+$$
+
+$$
+| \tilde{k}_i \rangle = \frac{\partial U_i}{\partial x_i} | k_i \rangle
+$$
+
+which allows us to expressed the expectation value derivative as:
+
+$$
+\text{Re}\big( \langle b_i | \tilde{k}_i \rangle \big)
+$$
+
+We can then also define $|k_i\rangle$ and $|b_i\rangle$ iteratively as:
+
+$$
+|b_{i-1} \rangle = U^{\dagger}_i |b_i\rangle \qquad |k\_{i-1}\rangle = U^{\dagger}\_{i-1} |k_i \rangle
+$$
+
+With the initial conditions:
+
+$$
+| b_n \rangle = 2  M |\Psi \rangle \qquad |k_n\rangle = U_{n-1} \dots |0\rangle
+$$
+
+where we iterate in reverse from $n$ to $0$.
+
+For the adjoint differentiation algorithm with expectation values, we iterate in reverse over $i$, taking an inner product at each step for each partial derivative.
+
+## Full statevector differentiation
+
+If we want to differentiate the entire statevector, we have:
+
+$$
+\frac{\partial}{\partial x_i} | \Psi \rangle = 
+U_n \dots \frac{\partial U_i}{\partial x_i} \dots U_0 |0\rangle
+$$
+
+At each step $i$, we have have the list of "Jacobian" statevectors.
+
+$$
+|J_{i,j}\rangle = U_i \dots \frac{\partial U_j}{\partial x_j} \dots U_0 |0\rangle
+$$
+
+Upon iteration, we update:
+
+$$
+| J_{i+1, j} \rangle = U_{i+1} |J_{i, j} \rangle
+$$
+
+And add a new element to the list
+
+$$
+| J_{i+1, i+1} \rangle = \frac{\partial U_{i+1}}{\partial x_{i+1}} U_{i} \dots U_0 |0\rangle = \frac{\partial U_{i+1}}{\partial x_{i+1}} | k_i \rangle
+$$
+
+We are able to make some improvements over performing one simulation per trainable parameter by iteratively calculating $|k_i\rangle$ and reusing it between each jacobian calculation.
+
+Note that this iterates from $0$ to $n$, instead if in reverse like expectation value differentiation.
+
+## Vector jacobian product
+
+When performing backpropagation across a workflow, we can reduce the computational complexity by calculating the <b>Vector Jacobian Product</b> (VJP) instead of the full jacobian product.
+
+For a full jacobian:
+
+$$
+J_{i, j} = \frac{\partial f_j(x)}{\partial x_i}
+$$
+
+And a workflow that depends on the output of $f$:
+
+$$
+y = h(f(x))
+$$
+
+we have the VJP:
+
+$$
+\frac{\partial y}{\partial x_i} = \frac{\partial y}{\partial f_j}  J_{i,j} 
+$$
+
+where we sum up over all $j$ indices.
+
+Calculating this quantity instead has several benefits over calculating the full jacobian.  
+
+1. We only need to calculate components for the full jacobian that actually influence final output.  If a component $j$  has $dy/df_j = 0$, we do not need to calculate derivatives for that output.
+
+2. Reduction in memory, which leads to a reduction in overall computational cost. The full jacobian scales with the size of the output *product* the size of the input.  But the VJP accepts something that is the size of the output and returns something that is the size of the input.
+
+For the VJP of expectation values, this looks like:
+
+$$
+\frac{\partial y}{\partial f_j}  \frac{\partial \langle M_j \rangle}{\partial x_i} = \frac{\partial \langle \Psi | \frac{\partial y}{\partial f_j} M_j |\Psi\rangle}{\partial x_i}
+$$
+
+So this is equivalent to performing a single jacobian calculation with a single "effective" observable of $\Sigma_j \frac{\partial y}{\partial f_j} M_j$.
+
+For the VJP calculation where the output is the state , we have:
+
+$$
+\frac{\partial y}{\partial f_j} \frac{\partial |\Psi\rangle}{\partial x_i}
+$$
+
+but we can make the identification that $\partial y/ \partial f_j$ is the same size and shape as a statevector, so we can simply rewrite this:
+
+$$
+\langle \left(\frac{\partial y}{\partial f} \right)^{\dagger} | \frac{\partial \Psi }{\partial x_i} \rangle
+$$
+
+$$
+\langle dy^{\dagger} | U_n \dots \frac{\partial U_i}{\partial x_i} \dots |0 \rangle
+$$ 
+
+We can then use the same algorithm as for adjoint differentiation but with the initial conditions 
+
+$$
+| b_n \rangle = |dy^{\dagger} \rangle  \qquad |k_n\rangle = U_{n-1} \dots |0\rangle
+$$
+
+instead of:
+
+$$
+| b_n \rangle = 2  M |\Psi \rangle \qquad |k_n\rangle = U_{n-1} \dots |0\rangle
+$$
\ No newline at end of file
diff --git a/pennylane/devices/qubit/adjoint_jacobian.py b/pennylane/devices/qubit/adjoint_jacobian.py
index 7258c78a0..d3ee8db1d 100644
--- a/pennylane/devices/qubit/adjoint_jacobian.py
+++ b/pennylane/devices/qubit/adjoint_jacobian.py
@@ -23,6 +23,7 @@ from pennylane.tape import QuantumTape
 
 from .apply_operation import apply_operation
 from .simulate import get_final_state
+from .initialize_state import create_initial_state
 
 # pylint: disable=protected-access, too-many-branches
 
@@ -34,6 +35,39 @@ def _dot_product_real(bra, ket, num_wires):
     return qml.math.real(qml.math.sum(qml.math.conj(bra) * ket, axis=sum_axes))
 
 
+def _adjoint_jacobian_state(tape: QuantumTape):
+    """Calculate the full jacobian for a circuit that returns the state.
+
+    Args:
+        tape (QuantumTape): the circuit we wish to differentiate
+
+    Returns:
+        TensorLike: the full jacobian.
+
+    See ``adjoint_jacobian.md`` for details on the algorithm.
+    """
+    jacobian = []
+
+    has_state_prep = isinstance(tape[0], qml.operation.StatePrepBase)
+    state = create_initial_state(tape.wires, tape[0] if has_state_prep else None)
+
+    param_idx = has_state_prep
+    for op in tape.operations[has_state_prep:]:
+        jacobian = [apply_operation(op, jac) for jac in jacobian]
+
+        if op.num_params == 1:
+            if param_idx in tape.trainable_params:
+                d_op_matrix = operation_derivative(op)
+                jacobian.append(
+                    apply_operation(qml.QubitUnitary(d_op_matrix, wires=op.wires), state)
+                )
+
+            param_idx += 1
+        state = apply_operation(op, state)
+
+    return tuple(jac.flatten() for jac in jacobian)
+
+
 def adjoint_jacobian(tape: QuantumTape, state=None):
     """Implements the adjoint method outlined in
     `Jones and Gacon <https://arxiv.org/abs/2009.02823>`__ to differentiate an input tape.
@@ -59,16 +93,20 @@ def adjoint_jacobian(tape: QuantumTape, state=None):
     Returns:
         array or tuple[array]: the derivative of the tape with respect to trainable parameters.
         Dimensions are ``(len(observables), len(trainable_params))``.
+
     """
     # Map wires if custom wire labels used
     tape = tape.map_to_standard_wires()
 
+    if isinstance(tape.measurements[0], qml.measurements.StateMP):
+        return _adjoint_jacobian_state(tape)
+
     ket = state if state is not None else get_final_state(tape)[0]
 
     n_obs = len(tape.observables)
     bras = np.empty([n_obs] + [2] * len(tape.wires), dtype=np.complex128)
     for kk, obs in enumerate(tape.observables):
-        bras[kk, ...] = apply_operation(obs, ket)
+        bras[kk, ...] = 2 * apply_operation(obs, ket)
 
     jac = np.zeros((len(tape.observables), len(tape.trainable_params)))
 
@@ -84,9 +122,7 @@ def adjoint_jacobian(tape: QuantumTape, state=None):
             if param_number in tape.trainable_params:
                 d_op_matrix = operation_derivative(op)
                 ket_temp = apply_operation(qml.QubitUnitary(d_op_matrix, wires=op.wires), ket)
-                jac[:, trainable_param_number] = 2 * _dot_product_real(
-                    bras, ket_temp, len(tape.wires)
-                )
+                jac[:, trainable_param_number] = _dot_product_real(bras, ket_temp, len(tape.wires))
 
                 trainable_param_number -= 1
             param_number -= 1
@@ -210,6 +246,8 @@ def adjoint_vjp(tape: QuantumTape, cotangents: Tuple[Number], state=None):
     Returns:
         Tuple[Number]: gradient vector for input parameters
     """
+    # See ``adjoint_jacobian.md`` to more information on the algorithm.
+
     # 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)}
@@ -218,29 +256,42 @@ def adjoint_vjp(tape: QuantumTape, cotangents: Tuple[Number], state=None):
 
     ket = state if state is not None else get_final_state(tape)[0]
 
-    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 = flat_bra.reshape(ket.shape)
+    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:
-        bra = apply_operation(obs, ket)
+        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)
 
     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))
+    cotangents_in = np.empty(len(tape.trainable_params), dtype=tape.measurements[0].numeric_type)
 
     for op in reversed(tape.operations[tape.num_preps :]):
         adj_op = qml.adjoint(op)
@@ -251,7 +302,9 @@ 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] = 2 * np.real(np.sum(np.conj(bra) * ket_temp))
+                cotangents_in[trainable_param_number] = real_if_expval(
+                    np.sum(np.conj(bra) * ket_temp)
+                )
 
                 trainable_param_number -= 1
             param_number -= 1
diff --git a/pennylane/interfaces/execution.py b/pennylane/interfaces/execution.py
index abdf2bb39..d6c05fbce 100644
--- a/pennylane/interfaces/execution.py
+++ b/pennylane/interfaces/execution.py
@@ -552,7 +552,6 @@ def execute(
         )
 
     ### Specifying and preprocessing variables ####
-    transform_program = transform_program or qml.transforms.core.TransformProgram()
 
     if interface == "auto":
         params = []
@@ -583,6 +582,12 @@ def execute(
         gradient_fn, grad_on_execution, interface, device, device_vjp
     )
 
+    if transform_program is None:
+        if isinstance(device, qml.devices.Device):
+            transform_program = device.preprocess(config)[0]
+        else:
+            transform_program = qml.transforms.core.TransformProgram()
+
     if isinstance(cache, bool) and cache:
         # cache=True: create a LRUCache object
         cache = LRUCache(maxsize=cachesize)
diff --git a/pennylane/interfaces/jacobian_products.py b/pennylane/interfaces/jacobian_products.py
index ab08c07b1..f405e99f2 100644
--- a/pennylane/interfaces/jacobian_products.py
+++ b/pennylane/interfaces/jacobian_products.py
@@ -57,7 +57,10 @@ def _compute_jvps(jacs, tangents, tapes):
         multi = len(t.measurements) > 1
         if len(t.trainable_params) == 0:
             empty_shots = qml.measurements.Shots(None)
-            zeros_jvp = tuple(np.zeros(mp.shape(None, empty_shots)) for mp in t.measurements)
+            zeros_jvp = tuple(
+                np.zeros(mp.shape(None, empty_shots), dtype=mp.numeric_type)
+                for mp in t.measurements
+            )
             zeros_jvp = zeros_jvp[0] if len(t.measurements) == 1 else zeros_jvp
             if t.shots.has_partitioned_shots:
                 jvps.append(tuple(zeros_jvp for _ in range(t.shots.num_copies)))
diff --git a/pennylane/interfaces/tensorflow.py b/pennylane/interfaces/tensorflow.py
index 2ae127c79..697b18a36 100644
--- a/pennylane/interfaces/tensorflow.py
+++ b/pennylane/interfaces/tensorflow.py
@@ -140,6 +140,12 @@ def _jac_restructured(jacs, tapes):
     return tuple(jacs_nested)
 
 
+def _recursive_conj(dy):
+    if isinstance(dy, (tf.Variable, tf.Tensor)):
+        return tf.math.conj(dy)
+    return tuple(_recursive_conj(d) for d in dy)
+
+
 def execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_diff=2):
     """Execute a batch of tapes with TensorFlow parameters on a device.
 
@@ -232,6 +238,10 @@ def execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_d
             # whether the tapes contain multiple measurements
             multi_measurements = [len(tape.measurements) > 1 for tape in tapes]
 
+            # TF obeys the dL/dz_conj convention instead of the
+            # dL/dz convention of PennyLane, autograd and jax. This converts between the formats
+            dy = _recursive_conj(dy)
+
             # reconstruct the nested structure of dy
             dy = _res_restructured(dy, tapes)
 
diff --git a/pennylane/interfaces/tensorflow_autograph.py b/pennylane/interfaces/tensorflow_autograph.py
index 8ad073b0d..0ec193540 100644
--- a/pennylane/interfaces/tensorflow_autograph.py
+++ b/pennylane/interfaces/tensorflow_autograph.py
@@ -180,6 +180,13 @@ def execute(
                 # No additional quantum evaluations needed; simply compute the VJPs directly.
 
                 def _backward(*args):
+                    for tape in tapes:
+                        for m in tape.measurements:
+                            if m.numeric_type == complex:
+                                raise NotImplementedError(
+                                    f"Tensorflow autograph only supports real valued measurements. Got {m}"
+                                )
+
                     dy = args[: total_measurements * num_shot_copies]
                     jacs = args[total_measurements * num_shot_copies : -len(tapes)]
                     multi_measurements = args[-len(tapes) :]
diff --git a/pennylane/interfaces/torch.py b/pennylane/interfaces/torch.py
index 66366903c..be10cddc8 100644
--- a/pennylane/interfaces/torch.py
+++ b/pennylane/interfaces/torch.py
@@ -109,6 +109,12 @@ def pytreeify(cls):
     return cls
 
 
+def _recursive_conj(dy):
+    if isinstance(dy, torch.Tensor):
+        return torch.conj(dy)
+    return tuple(_recursive_conj(d) for d in dy)
+
+
 @pytreeify
 class ExecuteTapes(torch.autograd.Function):
     """The signature of this ``torch.autograd.Function`` is designed to
@@ -175,9 +181,10 @@ class ExecuteTapes(torch.autograd.Function):
                     str(i) for i in inspect.getouterframes(inspect.currentframe(), 2)[1][1:3]
                 ),
             )
-
+        # Torch obeys the dL/dz_conj convention instead of the
+        # dL/dz convention of PennyLane, autograd and jax. This converts between the formats
+        dy = _recursive_conj(dy)
         vjps = ctx.jpc.compute_vjp(ctx.tapes, dy)
-
         # split tensor into separate entries
         unpacked_vjps = []
         for vjp_slice in vjps:
diff --git a/pennylane/measurements/state.py b/pennylane/measurements/state.py
index 4793213a4..8f21cae41 100644
--- a/pennylane/measurements/state.py
+++ b/pennylane/measurements/state.py
@@ -151,7 +151,7 @@ class StateMP(StateMeasurement):
         num_shot_elements = (
             sum(s.copies for s in shots.shot_vector) if shots.has_partitioned_shots else 1
         )
-        dim = 2 ** len(device.wires)
+        dim = 2 ** len(self.wires) if self.wires else 2 ** len(device.wires)
         return (dim,) if num_shot_elements == 1 else tuple((dim,) for _ in range(num_shot_elements))
 
     def process_state(self, state: Sequence[complex], wire_order: Wires):


Test patch:
diff --git a/tests/devices/default_qubit/test_default_qubit.py b/tests/devices/default_qubit/test_default_qubit.py
index 6d4f2cfe9..62292b4d9 100644
--- a/tests/devices/default_qubit/test_default_qubit.py
+++ b/tests/devices/default_qubit/test_default_qubit.py
@@ -132,7 +132,7 @@ class TestSupportsDerivatives:
         """Tests that DefaultQubit does not support adjoint differentiation with invalid circuits."""
         dev = DefaultQubit()
         config = ExecutionConfig(gradient_method="adjoint")
-        circuit = qml.tape.QuantumScript([], [qml.probs()])
+        circuit = qml.tape.QuantumScript([], [qml.sample()], shots=10)
         assert dev.supports_derivatives(config, circuit=circuit) is False
         assert dev.supports_jvp(config, circuit=circuit) is False
         assert dev.supports_vjp(config, circuit=circuit) is False
diff --git a/tests/devices/default_qubit/test_default_qubit_preprocessing.py b/tests/devices/default_qubit/test_default_qubit_preprocessing.py
index 4cbfcd470..e6dcfc05e 100644
--- a/tests/devices/default_qubit/test_default_qubit_preprocessing.py
+++ b/tests/devices/default_qubit/test_default_qubit_preprocessing.py
@@ -468,8 +468,8 @@ class TestPreprocessingIntegration:
         [
             (
                 [qml.RX(0.1, wires=0)],
-                [qml.probs(wires=[0, 1, 2])],
-                "not accepted for analytic simulation on adjoint",
+                [qml.probs(op=qml.PauliX(0))],
+                "adjoint diff supports either all expectation values or",
             ),
             (
                 [qml.RX(0.1, wires=0)],
@@ -642,7 +642,7 @@ class TestAdjointDiffTapeValidation:
         with pytest.raises(qml.DeviceError, match=msg):
             program((tape,))
 
-    def test_not_expval(self):
+    def test_non_diagonal_non_expval(self):
         """Test if a QuantumFunctionError is raised for a tape with measurements that are not
         expectation values"""
 
@@ -655,7 +655,7 @@ class TestAdjointDiffTapeValidation:
 
         with pytest.raises(
             qml.DeviceError,
-            match=r"not accepted for analytic simulation on adjoint",
+            match=r"adjoint diff supports either all expectation values or only measurements without observables",
         ):
             program((qs,))
 
diff --git a/tests/devices/qubit/test_adjoint_jacobian.py b/tests/devices/qubit/test_adjoint_jacobian.py
index 6bf0292e2..c78949ebe 100644
--- a/tests/devices/qubit/test_adjoint_jacobian.py
+++ b/tests/devices/qubit/test_adjoint_jacobian.py
@@ -311,6 +311,50 @@ class TestAdjointJacobian:
         assert np.allclose(grad_adjoint, expected)
 
 
+class TestAdjointJacobianState:
+    """Tests for differentiating a state vector."""
+
+    def test_simple_state_derivative(self):
+        """Test state differentiation for a single parameter."""
+        x = 1.2
+        tape = qml.tape.QuantumScript([qml.RX(x, wires=0)], [qml.state()])
+        jac = adjoint_jacobian(tape)
+        expected = [-0.5 * np.sin(x / 2), -0.5j * np.cos(x / 2)]
+        assert qml.math.allclose(jac, expected)
+
+        dy = np.array([0.5, 2.0])
+        vjp = adjoint_vjp(tape, dy)
+        expected_vjp = dy[0] * expected[0] + dy[1] * expected[1]
+        assert qml.math.allclose(expected_vjp, vjp)
+
+    def test_two_wires_two_parameters(self):
+        """Test a more complicated circuit with two parameters and two wires."""
+
+        x = 0.5
+        y = 0.6
+        tape = qml.tape.QuantumScript([qml.RX(x, 0), qml.RY(y, 1), qml.CNOT((0, 1))], [qml.state()])
+        x_jac, y_jac = adjoint_jacobian(tape)
+
+        c_x, s_x = np.cos(x / 2), np.sin(x / 2)
+        c_y, s_y = np.cos(y / 2), np.sin(y / 2)
+        x_jac_expected = np.array(
+            [-0.5 * c_y * s_x, -0.5 * s_y * s_x, -0.5j * s_y * c_x, -0.5j * c_x * c_y]
+        )
+        assert qml.math.allclose(x_jac, x_jac_expected)
+
+        y_jac_expected = np.array(
+            [-0.5 * c_x * s_y, 0.5 * c_x * c_y, -0.5j * s_x * c_y, 0.5j * s_x * s_y]
+        )
+        assert qml.math.allclose(y_jac, y_jac_expected)
+
+        dy = np.array([0.5, 1.0, 2.0, 2.5])
+        x_vjp, y_vjp = adjoint_vjp(tape, dy)
+        x_vjp_expected = np.dot(x_jac_expected, dy)
+        assert qml.math.allclose(x_vjp, x_vjp_expected)
+        y_vjp_expected = np.dot(y_jac_expected, dy)
+        assert qml.math.allclose(y_vjp, y_vjp_expected)
+
+
 class TestAdjointJVP:
     """Test for adjoint_jvp"""
 
diff --git a/tests/docs/test_supported_confs.py b/tests/docs/test_supported_confs.py
index 0e43bbe77..911313a74 100644
--- a/tests/docs/test_supported_confs.py
+++ b/tests/docs/test_supported_confs.py
@@ -37,6 +37,7 @@ from pennylane.measurements import (
     VnEntropy,
     MutualInfo,
 )
+from pennylane.measurements.measurements import ObservableReturnTypes
 
 pytestmark = pytest.mark.all_interfaces
 
@@ -162,7 +163,7 @@ def get_variable(interface, wire_specs, complex=False):
         return np.array([0.1] * num_wires, requires_grad=True)
     if interface == "jax":
         # complex dtype is required for JAX when holomorphic gradient is used
-        return jnp.array([0.1] * num_wires, dtype=np.complex64 if complex else None)
+        return jnp.array([0.1] * num_wires, dtype=np.complex128 if complex else None)
     if interface == "tf":
         # complex dtype is required for TF when the gradients have non-zero
         # imaginary parts, otherwise they will be ignored
@@ -308,8 +309,6 @@ class TestSupportedConfs:
                     msg = "not accepted with finite shots"
                 else:
                     msg = "Finite shots are not supported with"
-            elif return_type not in ("Hermitian", "Projector", Expectation):
-                msg = "not accepted for analytic simulation"
         elif shots and return_type in (
             VnEntropy,
             MutualInfo,
@@ -366,15 +365,19 @@ class TestSupportedConfs:
     def test_all_adjoint_nonexp(self, interface, return_type, shots, wire_specs):
         """Test diff_method=adjoint raises an error for non-expectation
         measurements for all interfaces"""
-        msg = (
-            ""  # two options of message both along the lines of "no shots"
-            if shots
-            else "not accepted for analytic simulation on adjoint"
-        )
 
         circuit = get_qnode(interface, "adjoint", return_type, shots, wire_specs)
-        x = get_variable(interface, wire_specs)
-        with pytest.raises(qml.DeviceError, match=msg):
+        x = get_variable(interface, wire_specs, complex=True)
+
+        if shots is not None:
+            with pytest.raises(qml.DeviceError):
+                compute_gradient(x, interface, circuit, return_type)
+        elif return_type == ObservableReturnTypes.Probability and interface == "tf":
+            with pytest.raises(Exception):
+                # tensorflow.python.framework.errors_impl.InvalidArgumentError
+                # TODO: figure out why
+                compute_gradient(x, interface, circuit, return_type)
+        else:
             compute_gradient(x, interface, circuit, return_type)
 
     @pytest.mark.parametrize("interface", diff_interfaces)
diff --git a/tests/interfaces/default_qubit_2_integration/test_autograd_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_autograd_default_qubit_2.py
index 31a900f6a..18d8035e0 100644
--- a/tests/interfaces/default_qubit_2_integration/test_autograd_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_autograd_default_qubit_2.py
@@ -266,9 +266,6 @@ class TestAutogradExecuteIntegration:
         """Test that a tape with no parameters is correctly
         ignored during the gradient computation"""
 
-        if execute_kwargs["gradient_fn"] == "adjoint":
-            pytest.skip("Adjoint differentiation does not yet support probabilities")
-
         def cost(params):
             tape1 = qml.tape.QuantumScript(
                 [qml.Hadamard(0)], [qml.expval(qml.PauliX(0))], shots=shots
@@ -567,9 +564,6 @@ class TestAutogradExecuteIntegration:
         """Tests correct output shape and evaluation for a tape
         with prob outputs"""
 
-        if execute_kwargs["gradient_fn"] == "adjoint":
-            pytest.skip("adjoint differentiation does not suppport probabilities.")
-
         def cost(x, y):
             ops = [qml.RX(x, 0), qml.RY(y, 1), qml.CNOT((0, 1))]
             m = [qml.probs(wires=0), qml.probs(wires=1)]
@@ -622,8 +616,6 @@ class TestAutogradExecuteIntegration:
     def test_ragged_differentiation(self, execute_kwargs, device, shots):
         """Tests correct output shape and evaluation for a tape
         with prob and expval outputs"""
-        if execute_kwargs["gradient_fn"] == "adjoint":
-            pytest.skip("Adjoint differentiation does not yet support probabilities")
 
         def cost(x, y):
             ops = [qml.RX(x, wires=0), qml.RY(y, 1), qml.CNOT((0, 1))]
diff --git a/tests/interfaces/default_qubit_2_integration/test_autograd_qnode_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_autograd_qnode_default_qubit_2.py
index 94dc7627c..412498e4d 100644
--- a/tests/interfaces/default_qubit_2_integration/test_autograd_qnode_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_autograd_qnode_default_qubit_2.py
@@ -560,9 +560,7 @@ class TestQubitIntegration:
             device_vjp=device_vjp,
         )
 
-        if diff_method == "adjoint":
-            pytest.skip("The adjoint method does not currently support returning probabilities")
-        elif diff_method == "spsa":
+        if diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
             tol = TOL_FOR_SPSA
 
@@ -597,9 +595,7 @@ class TestQubitIntegration:
             device_vjp=device_vjp,
         )
 
-        if diff_method == "adjoint":
-            pytest.skip("The adjoint method does not currently support returning probabilities")
-        elif diff_method == "spsa":
+        if diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
             tol = TOL_FOR_SPSA
 
@@ -663,9 +659,7 @@ class TestQubitIntegration:
             device_vjp=device_vjp,
         )
 
-        if diff_method == "adjoint":
-            pytest.skip("The adjoint method does not currently support returning probabilities")
-        elif diff_method == "spsa":
+        if diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
             tol = TOL_FOR_SPSA
 
@@ -714,9 +708,7 @@ class TestQubitIntegration:
             device_vjp=device_vjp,
         )
 
-        if diff_method == "adjoint":
-            pytest.skip("The adjoint method does not currently support returning probabilities")
-        elif diff_method == "spsa":
+        if diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
             tol = TOL_FOR_SPSA
         elif diff_method == "hadamard":
@@ -1306,8 +1298,6 @@ class TestQubitIntegration:
 
     def test_state(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol):
         """Test that the state can be returned and differentiated"""
-        if diff_method == "adjoint":
-            pytest.skip("Adjoint does not support states")
 
         x = np.array(0.543, requires_grad=True)
         y = np.array(-0.654, requires_grad=True)
@@ -1352,6 +1342,8 @@ class TestQubitIntegration:
         self, state, interface, dev, diff_method, grad_on_execution, device_vjp, tol
     ):
         """Test that the variance of a projector is correctly returned"""
+        if diff_method == "adjoint":
+            pytest.skip("adjoint supports either expvals or diagonal measurements.")
         kwargs = dict(
             diff_method=diff_method,
             interface=interface,
@@ -1359,9 +1351,7 @@ class TestQubitIntegration:
             device_vjp=device_vjp,
         )
 
-        if diff_method == "adjoint":
-            pytest.skip("Adjoint does not support projectors")
-        elif diff_method == "spsa":
+        if diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
             tol = TOL_FOR_SPSA
         elif diff_method == "hadamard":
@@ -1824,8 +1814,6 @@ class TestReturn:
     ):
         """For a multi dimensional measurement (probs), check that a single array is returned with the correct
         dimension"""
-        if diff_method == "adjoint":
-            pytest.skip("The adjoint method does not currently support returning probabilities")
 
         @qnode(
             dev,
@@ -1851,8 +1839,6 @@ class TestReturn:
     ):
         """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with
         the correct dimension"""
-        if diff_method == "adjoint":
-            pytest.skip("The adjoint method does not currently support returning probabilities")
 
         @qnode(
             dev,
@@ -1883,8 +1869,6 @@ class TestReturn:
         self, dev, diff_method, grad_on_execution, device_vjp
     ):
         """For a multi dimensional measurement (probs), check that a single array is returned."""
-        if diff_method == "adjoint":
-            pytest.skip("The adjoint method does not currently support returning probabilities")
 
         @qnode(
             dev,
@@ -1909,9 +1893,6 @@ class TestReturn:
     ):
         """The jacobian of multiple measurements with a single params return an array."""
 
-        if diff_method == "adjoint":
-            pytest.skip("The adjoint method does not currently support returning probabilities")
-
         @qnode(
             dev,
             interface="autograd",
@@ -1939,9 +1920,6 @@ class TestReturn:
     ):
         """The jacobian of multiple measurements with a multiple params return a tuple of arrays."""
 
-        if diff_method == "adjoint":
-            pytest.skip("The adjoint method does not currently support returning probabilities")
-
         @qnode(
             dev,
             interface="autograd",
@@ -1976,9 +1954,6 @@ class TestReturn:
     ):
         """The jacobian of multiple measurements with a multiple params array return a single array."""
 
-        if diff_method == "adjoint":
-            pytest.skip("The adjoint method does not currently support returning probabilities")
-
         @qnode(
             dev,
             interface="autograd",
diff --git a/tests/interfaces/default_qubit_2_integration/test_execute_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_execute_default_qubit_2.py
index 725784671..618620e34 100644
--- a/tests/interfaces/default_qubit_2_integration/test_execute_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_execute_default_qubit_2.py
@@ -161,8 +161,8 @@ def test_caching(gradient_fn):
     assert len(cache) == 1
     assert cache[qs.hash] == -1.0
 
-    assert results == [-1.0, -1.0]
-    assert results2 == [-1.0, -1.0]
+    assert list(results) == [-1.0, -1.0]
+    assert list(results2) == [-1.0, -1.0]
 
     assert tracker.totals["batches"] == 1
     assert tracker.totals["executions"] == 1
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 edd88e4df..302f42325 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
@@ -245,9 +245,6 @@ class TestJaxExecuteIntegration:
         """Test that a tape with no parameters is correctly
         ignored during the gradient computation"""
 
-        if execute_kwargs["gradient_fn"] == "adjoint":
-            pytest.skip("Adjoint differentiation does not yet support probabilities")
-
         def cost(params):
             tape1 = qml.tape.QuantumScript(
                 [qml.Hadamard(0)], [qml.expval(qml.PauliX(0))], shots=shots
@@ -276,6 +273,7 @@ class TestJaxExecuteIntegration:
             return out
 
         params = jnp.array([0.1, 0.2])
+
         x, y = params
 
         res = cost(params)
@@ -539,9 +537,6 @@ class TestJaxExecuteIntegration:
         """Tests correct output shape and evaluation for a tape
         with prob outputs"""
 
-        if execute_kwargs["gradient_fn"] == "adjoint":
-            pytest.skip("adjoint differentiation does not support probabilities")
-
         def cost(x, y):
             ops = [qml.RX(x, 0), qml.RY(y, 1), qml.CNOT((0, 1))]
             m = [qml.probs(wires=0), qml.probs(wires=1)]
@@ -569,6 +564,10 @@ 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,)
@@ -604,8 +603,6 @@ class TestJaxExecuteIntegration:
     def test_ragged_differentiation(self, execute_kwargs, device, shots):
         """Tests correct output shape and evaluation for a tape
         with prob and expval outputs"""
-        if execute_kwargs["gradient_fn"] == "adjoint":
-            pytest.skip("Adjoint differentiation does not yet support probabilities")
 
         def cost(x, y):
             ops = [qml.RX(x, wires=0), qml.RY(y, 1), qml.CNOT((0, 1))]
@@ -628,6 +625,10 @@ 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,)
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 43af86e31..1f74286ca 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
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Integration tests for using the JAX-Python interface with a QNode"""
-# pylint: disable=no-member, too-many-arguments, unexpected-keyword-arg
+# pylint: disable=no-member, too-many-arguments, unexpected-keyword-arg, use-implicit-booleaness-not-comparison
 
 from itertools import product
 import numpy as np
@@ -187,12 +187,12 @@ class TestQNode:
     ):
         """Test that operation and nested tape expansion
         is differentiable"""
-        kwargs = dict(
-            diff_method=diff_method,
-            interface=interface,
-            grad_on_execution=grad_on_execution,
-            device_vjp=device_vjp,
-        )
+        kwargs = {
+            "diff_method": diff_method,
+            "interface": interface,
+            "grad_on_execution": grad_on_execution,
+            "device_vjp": device_vjp,
+        }
 
         if diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
@@ -265,12 +265,12 @@ class TestVectorValuedQNode:
         self, dev, diff_method, grad_on_execution, device_vjp, interface, tol
     ):
         """Test jacobian calculation"""
-        kwargs = dict(
-            diff_method=diff_method,
-            interface=interface,
-            grad_on_execution=grad_on_execution,
-            device_vjp=device_vjp,
-        )
+        kwargs = {
+            "diff_method": diff_method,
+            "interface": interface,
+            "grad_on_execution": grad_on_execution,
+            "device_vjp": device_vjp,
+        }
 
         if diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
@@ -326,12 +326,12 @@ class TestVectorValuedQNode:
         self, dev, diff_method, grad_on_execution, device_vjp, interface, tol
     ):
         """Test jacobian calculation when no prior circuit evaluation has been performed"""
-        kwargs = dict(
-            diff_method=diff_method,
-            interface=interface,
-            grad_on_execution=grad_on_execution,
-            device_vjp=device_vjp,
-        )
+        kwargs = {
+            "diff_method": diff_method,
+            "interface": interface,
+            "grad_on_execution": grad_on_execution,
+            "device_vjp": device_vjp,
+        }
 
         if diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
@@ -384,15 +384,13 @@ class TestVectorValuedQNode:
     ):
         """Tests correct output shape and evaluation for a tape
         with a single prob output"""
-        kwargs = dict(
-            diff_method=diff_method,
-            interface=interface,
-            grad_on_execution=grad_on_execution,
-            device_vjp=device_vjp,
-        )
-        if diff_method == "adjoint":
-            pytest.skip("Adjoint does not support probs")
-        elif diff_method == "spsa":
+        kwargs = {
+            "diff_method": diff_method,
+            "interface": interface,
+            "grad_on_execution": grad_on_execution,
+            "device_vjp": device_vjp,
+        }
+        if diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
             tol = TOL_FOR_SPSA
 
@@ -406,6 +404,10 @@ 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(
@@ -432,16 +434,14 @@ class TestVectorValuedQNode:
     ):
         """Tests correct output shape and evaluation for a tape
         with multiple prob outputs"""
-        kwargs = dict(
-            diff_method=diff_method,
-            interface=interface,
-            grad_on_execution=grad_on_execution,
-            device_vjp=device_vjp,
-        )
+        kwargs = {
+            "diff_method": diff_method,
+            "interface": interface,
+            "grad_on_execution": grad_on_execution,
+            "device_vjp": device_vjp,
+        }
 
-        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
 
@@ -473,6 +473,10 @@ 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(
             [
@@ -514,15 +518,13 @@ class TestVectorValuedQNode:
     ):
         """Tests correct output shape and evaluation for a tape
         with prob and expval outputs"""
-        kwargs = dict(
-            diff_method=diff_method,
-            interface=interface,
-            grad_on_execution=grad_on_execution,
-            device_vjp=device_vjp,
-        )
-        if diff_method == "adjoint":
-            pytest.skip("Adjoint does not support probs")
-        elif diff_method == "spsa":
+        kwargs = {
+            "diff_method": diff_method,
+            "interface": interface,
+            "grad_on_execution": grad_on_execution,
+            "device_vjp": device_vjp,
+        }
+        if diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
             tol = TOL_FOR_SPSA
 
@@ -549,6 +551,10 @@ 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],
@@ -585,15 +591,17 @@ class TestVectorValuedQNode:
         """Tests correct output shape and evaluation for a tape with prob and expval outputs with less
         trainable parameters (argnums) than parameters."""
         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
 
         x = jax.numpy.array(0.543)
         y = jax.numpy.array(-0.654)
 
+        if diff_method == "adjoint":
+            x = x + 0j
+            y = y + 0j
+
         @qnode(
             dev,
             diff_method=diff_method,
@@ -608,6 +616,10 @@ 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 = [
@@ -635,16 +647,14 @@ class TestVectorValuedQNode:
     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"""
-        kwargs = dict(
-            diff_method=diff_method,
-            interface=interface,
-            grad_on_execution=grad_on_execution,
-            device_vjp=device_vjp,
-        )
-
-        if diff_method == "adjoint":
-            pytest.skip("Adjoint does not support probs")
-        elif diff_method == "hadamard":
+        kwargs = {
+            "diff_method": diff_method,
+            "interface": interface,
+            "grad_on_execution": grad_on_execution,
+            "device_vjp": device_vjp,
+        }
+
+        if diff_method == "hadamard":
             pytest.skip("Hadamard does not support var")
         elif diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
@@ -675,6 +685,10 @@ 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],
@@ -945,13 +959,13 @@ class TestQubitIntegrationHigherOrder:
         self, dev, diff_method, grad_on_execution, device_vjp, interface, tol
     ):
         """Test second derivative calculation of a scalar-valued QNode"""
-        kwargs = dict(
-            diff_method=diff_method,
-            interface=interface,
-            grad_on_execution=grad_on_execution,
-            device_vjp=device_vjp,
-            max_diff=2,
-        )
+        kwargs = {
+            "diff_method": diff_method,
+            "interface": interface,
+            "grad_on_execution": grad_on_execution,
+            "device_vjp": device_vjp,
+            "max_diff": 2,
+        }
 
         if diff_method == "adjoint":
             pytest.skip("Adjoint does not second derivative.")
@@ -1237,8 +1251,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)
@@ -1278,8 +1290,8 @@ class TestQubitIntegrationHigherOrder:
         """Test that the variance of a projector is correctly returned"""
         gradient_kwargs = {}
         if diff_method == "adjoint":
-            pytest.skip("Adjoint does not support projectors")
-        elif diff_method == "hadamard":
+            pytest.skip("adjoint supports all expvals or only diagonal measurements.")
+        if diff_method == "hadamard":
             pytest.skip("Hadamard does not support var.")
         elif diff_method == "spsa":
             gradient_kwargs = {"h": H_FOR_SPSA, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)}
@@ -1794,8 +1806,8 @@ class TestReturn:  # pylint:disable=too-many-public-methods
     ):
         """The hessian of multiple measurements with multiple params return a tuple of arrays."""
         if diff_method == "adjoint":
-            pytest.skip("Test does not supports adjoint because of var.")
-        elif diff_method == "hadamard":
+            pytest.skip("adjoint supports either all measurements or only diagonal measurements.")
+        if diff_method == "hadamard":
             pytest.skip("Test does not support 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")
@@ -1842,8 +1854,8 @@ class TestReturn:  # pylint:disable=too-many-public-methods
     ):
         """The jacobian of multiple measurements with a multiple params array return a single array."""
         if diff_method == "adjoint":
-            pytest.skip("Test does not supports adjoint because of var.")
-        elif diff_method == "hadamard":
+            pytest.skip("adjoint supports either all expvals or all diagonal measurements.")
+        if diff_method == "hadamard":
             pytest.skip("Test does not support 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")
@@ -1881,8 +1893,8 @@ class TestReturn:  # pylint:disable=too-many-public-methods
         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 diff_method == "adjoint" and jacobian == jax.jacfwd:
+            pytest.skip("jacfwd doesn't like complex numbers")
 
         @qnode(
             dev,
@@ -1898,6 +1910,10 @@ 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)
@@ -1914,11 +1930,12 @@ class TestReturn:  # pylint:disable=too-many-public-methods
         self, dev, diff_method, grad_on_execution, jacobian, device_vjp, shots, interface
     ):
         """The jacobian of multiple measurements with a multiple params return a tuple of arrays."""
-        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 diff_method == "adjoint" and jacobian == jax.jacfwd:
+            pytest.skip("jacfwd doesn't like complex numbers")
+
         @qnode(
             dev,
             interface=interface,
@@ -1934,6 +1951,10 @@ 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)
@@ -1958,11 +1979,12 @@ class TestReturn:  # pylint:disable=too-many-public-methods
         self, dev, diff_method, grad_on_execution, jacobian, device_vjp, shots, interface
     ):
         """The jacobian of multiple measurements with a multiple params array return a single array."""
-        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 diff_method == "adjoint" and jacobian == jax.jacfwd:
+            pytest.skip("jacfwd doesn't like complex numbers")
+
         @qnode(
             dev,
             interface=interface,
@@ -1977,6 +1999,10 @@ 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/default_qubit_2_integration/test_tensorflow_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_tensorflow_default_qubit_2.py
index 5e3eb211b..fc4084703 100644
--- a/tests/interfaces/default_qubit_2_integration/test_tensorflow_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_tensorflow_default_qubit_2.py
@@ -109,14 +109,14 @@ class TestCaching:
 # add tests for lightning 2 when possible
 # set rng for device when possible
 test_matrix = [
-    ({"gradient_fn": param_shift, "interface": "tensorflow"}, 100000, DefaultQubit(seed=42)),
-    ({"gradient_fn": param_shift, "interface": "tensorflow"}, None, DefaultQubit()),
-    ({"gradient_fn": "backprop", "interface": "tensorflow"}, None, DefaultQubit()),
-    ({"gradient_fn": "adjoint", "interface": "tensorflow"}, None, DefaultQubit()),
-    ({"gradient_fn": param_shift, "interface": "tf-autograph"}, 100000, DefaultQubit(seed=42)),
-    ({"gradient_fn": param_shift, "interface": "tf-autograph"}, None, DefaultQubit()),
-    ({"gradient_fn": "backprop", "interface": "tf-autograph"}, None, DefaultQubit()),
-    ({"gradient_fn": "adjoint", "interface": "tf-autograph"}, None, DefaultQubit()),
+    ({"gradient_fn": param_shift, "interface": "tensorflow"}, 100000, DefaultQubit(seed=42)),  # 0
+    ({"gradient_fn": param_shift, "interface": "tensorflow"}, None, DefaultQubit()),  # 1
+    ({"gradient_fn": "backprop", "interface": "tensorflow"}, None, DefaultQubit()),  # 2
+    ({"gradient_fn": "adjoint", "interface": "tensorflow"}, None, DefaultQubit()),  # 3
+    ({"gradient_fn": param_shift, "interface": "tf-autograph"}, 100000, DefaultQubit(seed=42)),  # 4
+    ({"gradient_fn": param_shift, "interface": "tf-autograph"}, None, DefaultQubit()),  # 5
+    ({"gradient_fn": "backprop", "interface": "tf-autograph"}, None, DefaultQubit()),  # 6
+    ({"gradient_fn": "adjoint", "interface": "tf-autograph"}, None, DefaultQubit()),  # 7
 ]
 
 
@@ -212,9 +212,6 @@ class TestTensorflowExecuteIntegration:
         """Test that a tape with no parameters is correctly
         ignored during the gradient computation"""
 
-        if execute_kwargs["gradient_fn"] == "adjoint":
-            pytest.skip("Adjoint differentiation does not yet support probabilities")
-
         def cost(params):
             tape1 = qml.tape.QuantumScript(
                 [qml.Hadamard(0)], [qml.expval(qml.PauliX(0))], shots=shots
@@ -252,6 +249,14 @@ class TestTensorflowExecuteIntegration:
         expected = 2 + tf.cos(0.5) + tf.cos(x) * tf.cos(y)
         assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0)
 
+        if (
+            execute_kwargs.get("interface", "") == "tf-autograph"
+            and execute_kwargs.get("gradient_fn", "") == "adjoint"
+        ):
+            with pytest.raises(NotImplementedError):
+                tape.gradient(res, params)
+            return
+
         grad = tape.gradient(res, params)
         expected = [-tf.cos(y) * tf.sin(x), -tf.cos(x) * tf.sin(y)]
         assert np.allclose(grad, expected, atol=atol_for_shots(shots), rtol=0)
@@ -499,9 +504,6 @@ class TestTensorflowExecuteIntegration:
         """Tests correct output shape and evaluation for a tape
         with prob outputs"""
 
-        if execute_kwargs["gradient_fn"] == "adjoint":
-            pytest.skip("adjoint differentiation does not support probabilities")
-
         def cost(x, y):
             ops = [qml.RX(x, 0), qml.RY(y, 1), qml.CNOT((0, 1))]
             m = [qml.probs(wires=0), qml.probs(wires=1)]
@@ -526,6 +528,13 @@ class TestTensorflowExecuteIntegration:
         )
         assert np.allclose(cost_res, expected, atol=atol_for_shots(shots), rtol=0)
 
+        if (
+            execute_kwargs.get("interface", "") == "tf-autograph"
+            and execute_kwargs.get("gradient_fn", "") == "adjoint"
+        ):
+            with pytest.raises(tf.errors.UnimplementedError):
+                tape.jacobian(cost_res, [x, y])
+            return
         res = tape.jacobian(cost_res, [x, y])
         assert isinstance(res, list) and len(res) == 2
         assert res[0].shape == (4,)
@@ -555,8 +564,6 @@ class TestTensorflowExecuteIntegration:
     def test_ragged_differentiation(self, execute_kwargs, device, shots):
         """Tests correct output shape and evaluation for a tape
         with prob and expval outputs"""
-        if execute_kwargs["gradient_fn"] == "adjoint":
-            pytest.skip("Adjoint differentiation does not yet support probabilities")
 
         def cost(x, y):
             ops = [qml.RX(x, wires=0), qml.RY(y, 1), qml.CNOT((0, 1))]
@@ -575,6 +582,13 @@ class TestTensorflowExecuteIntegration:
         )
         assert np.allclose(cost_res, expected, atol=atol_for_shots(shots), rtol=0)
 
+        if (
+            execute_kwargs.get("interface", "") == "tf-autograph"
+            and execute_kwargs.get("gradient_fn", "") == "adjoint"
+        ):
+            with pytest.raises(tf.errors.UnimplementedError):
+                tape.jacobian(cost_res, [x, y])
+            return
         res = tape.jacobian(cost_res, [x, y])
         assert isinstance(res, list) and len(res) == 2
         assert res[0].shape == (3,)
diff --git a/tests/interfaces/default_qubit_2_integration/test_tensorflow_qnode_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_tensorflow_qnode_default_qubit_2.py
index 18b045f59..7a4a6d367 100644
--- a/tests/interfaces/default_qubit_2_integration/test_tensorflow_qnode_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_tensorflow_qnode_default_qubit_2.py
@@ -512,9 +512,7 @@ class TestQubitIntegration:
         kwargs = dict(
             diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface
         )
-        if diff_method == "adjoint":
-            pytest.skip("The adjoint method does not currently support returning probabilities")
-        elif diff_method == "spsa":
+        if diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
             kwargs["num_directions"] = 20
             tol = TOL_FOR_SPSA
@@ -562,9 +560,7 @@ class TestQubitIntegration:
         kwargs = dict(
             diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface
         )
-        if diff_method == "adjoint":
-            pytest.skip("The adjoint method does not currently support returning probabilities")
-        elif diff_method == "spsa":
+        if diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
             kwargs["num_directions"] = 20
             tol = TOL_FOR_SPSA
@@ -846,8 +842,6 @@ class TestQubitIntegration:
 
     def test_state(self, dev, diff_method, grad_on_execution, tol, interface):
         """Test that the state can be returned and differentiated"""
-        if diff_method == "adjoint":
-            pytest.skip("Adjoint does not support states")
 
         x = tf.Variable(0.543, dtype=tf.float64)
         y = tf.Variable(-0.654, dtype=tf.float64)
@@ -883,8 +877,8 @@ class TestQubitIntegration:
             diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface
         )
         if diff_method == "adjoint":
-            pytest.skip("Adjoint does not support projectors")
-        elif diff_method == "hadamard":
+            pytest.skip("adjoint supports either all expvals or all diagonal measurements.")
+        if diff_method == "hadamard":
             pytest.skip("Variance not implemented yet.")
         elif diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
@@ -1628,8 +1622,6 @@ class TestReturn:
     ):
         """For a multi dimensional measurement (probs), check that a single array is returned with the correct
         dimension"""
-        if diff_method == "adjoint":
-            pytest.skip("Test does not supports adjoint because of probabilities.")
 
         @qnode(
             dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
@@ -1654,8 +1646,6 @@ class TestReturn:
     ):
         """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with
         the correct dimension"""
-        if diff_method == "adjoint":
-            pytest.skip("Test does not supports adjoint because of probabilities.")
 
         @qnode(
             dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
@@ -1685,8 +1675,6 @@ class TestReturn:
         self, dev, diff_method, grad_on_execution, interface
     ):
         """For a multi dimensional measurement (probs), check that a single array is returned."""
-        if diff_method == "adjoint":
-            pytest.skip("Test does not supports adjoint because of probabilities.")
 
         @qnode(
             dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
@@ -1710,8 +1698,6 @@ class TestReturn:
         self, dev, diff_method, grad_on_execution, interface
     ):
         """The jacobian of multiple measurements with a single params return an array."""
-        if diff_method == "adjoint":
-            pytest.skip("Test does not supports adjoint because of probabilities.")
 
         @qnode(
             dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
@@ -1737,9 +1723,6 @@ class TestReturn:
     ):
         """The jacobian of multiple measurements with a multiple params return a tuple of arrays."""
 
-        if diff_method == "adjoint":
-            pytest.skip("Test does not supports adjoint because of probabilities.")
-
         @qnode(
             dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
         )
@@ -1771,9 +1754,6 @@ class TestReturn:
     ):
         """The jacobian of multiple measurements with a multiple params array return a single array."""
 
-        if diff_method == "adjoint":
-            pytest.skip("Test does not supports adjoint because of probabilities.")
-
         @qnode(
             dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
         )
diff --git a/tests/interfaces/default_qubit_2_integration/test_torch_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_torch_default_qubit_2.py
index b6dff017a..e98aaf15d 100644
--- a/tests/interfaces/default_qubit_2_integration/test_torch_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_torch_default_qubit_2.py
@@ -285,9 +285,6 @@ class TestTorchExecuteIntegration:
         """Test that a tape with no parameters is correctly
         ignored during the gradient computation"""
 
-        if execute_kwargs["gradient_fn"] == "adjoint":
-            pytest.skip("Adjoint differentiation does not yet support probabilities")
-
         def cost(params):
             tape1 = qml.tape.QuantumScript(
                 [qml.Hadamard(0)], [qml.expval(qml.PauliX(0))], shots=shots
@@ -566,9 +563,6 @@ class TestTorchExecuteIntegration:
         """Tests correct output shape and evaluation for a tape
         with prob outputs"""
 
-        if execute_kwargs["gradient_fn"] == "adjoint":
-            pytest.skip("adjoint differentiation does not suppport probabilities.")
-
         def cost(x, y):
             ops = [qml.RX(x, 0), qml.RY(y, 1), qml.CNOT((0, 1))]
             m = [qml.probs(wires=0), qml.probs(wires=1)]
@@ -620,8 +614,6 @@ class TestTorchExecuteIntegration:
     def test_ragged_differentiation(self, execute_kwargs, device, shots):
         """Tests correct output shape and evaluation for a tape
         with prob and expval outputs"""
-        if execute_kwargs["gradient_fn"] == "adjoint":
-            pytest.skip("Adjoint differentiation does not yet support probabilities")
 
         def cost(x, y):
             ops = [qml.RX(x, wires=0), qml.RY(y, 1), qml.CNOT((0, 1))]
diff --git a/tests/interfaces/default_qubit_2_integration/test_torch_qnode_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_torch_qnode_default_qubit_2.py
index 39f171e71..ad91bb2b3 100644
--- a/tests/interfaces/default_qubit_2_integration/test_torch_qnode_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_torch_qnode_default_qubit_2.py
@@ -651,9 +651,7 @@ class TestQubitIntegration:
         """Tests correct output shape and evaluation for a tape
         with prob and expval outputs"""
         kwargs = {}
-        if diff_method == "adjoint":
-            pytest.skip("The adjoint method does not currently support returning probabilities")
-        elif diff_method == "spsa":
+        if diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
             tol = TOL_FOR_SPSA
 
@@ -714,9 +712,7 @@ class TestQubitIntegration:
             interface=interface,
             device_vjp=device_vjp,
         )
-        if diff_method == "adjoint":
-            pytest.skip("The adjoint method does not currently support returning probabilities")
-        elif diff_method == "spsa":
+        if diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
             kwargs["num_directions"] = 20
             tol = TOL_FOR_SPSA
@@ -1057,8 +1053,6 @@ class TestQubitIntegration:
 
     def test_state(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol):
         """Test that the state can be returned and differentiated"""
-        if diff_method == "adjoint":
-            pytest.skip("Adjoint does not support states")
 
         x = torch.tensor(0.543, requires_grad=True)
         y = torch.tensor(-0.654, requires_grad=True)
@@ -1106,8 +1100,8 @@ class TestQubitIntegration:
             device_vjp=device_vjp,
         )
         if diff_method == "adjoint":
-            pytest.skip("Adjoint does not support projectors")
-        elif diff_method == "spsa":
+            pytest.skip("adjoint supports either all expvals or all diagonal measurements")
+        if diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
             kwargs["num_directions"] = 20
             tol = TOL_FOR_SPSA