Repository: PennyLaneAI/pennylane
Pull request #3790 (2023-02-17 19:59:49)

Add function for adjoint differentiation to `devices.qubit`
**Context:**
Adding adjoint differentiation support for `DefaultQubit2`.

**Description of the Change:**
* Added `adjoint_jacobian` function inside `devices/qubit/adjoint_jacobian.py`

**Benefits:**
Adjoint differentiation can now be performed on the new qubit state-vector device.


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://codecov.io/gh/PennyLaneAI/pennylane/pull/3790?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3790](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3790?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (79ffc91) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/1774fe36f31f7de6266cf44983c85dfb561ec0fc?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (1774fe3) will **decrease** coverage by `0.18%`.
> The diff coverage is `0.00%`.

```diff
@@            Coverage Diff             @@
##           master    #3790      +/-   ##
==========================================
- Coverage   99.82%   99.64%   -0.18%     
==========================================
  Files         330      335       +5     
  Lines       28943    29422     +479     
==========================================
+ Hits        28892    29319     +427     
- Misses         51      103      +52     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3790?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/devices/qubit/adjoint\_jacobian.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3790?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvcXViaXQvYWRqb2ludF9qYWNvYmlhbi5weQ==) | `0.00% <0.00%> (ø)` | |
| [pennylane/ops/op\_math/evolution.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3790?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9vcF9tYXRoL2V2b2x1dGlvbi5weQ==) | `97.50% <0.00%> (-2.50%)` | :arrow_down: |
| [pennylane/\_qutrit\_device.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3790?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19xdXRyaXRfZGV2aWNlLnB5) | `98.07% <0.00%> (-1.93%)` | :arrow_down: |
| [pennylane/numpy/tensor.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3790?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL251bXB5L3RlbnNvci5weQ==) | `98.80% <0.00%> (-1.20%)` | :arrow_down: |
| [pennylane/\_qubit\_device.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3790?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19xdWJpdF9kZXZpY2UucHk=) | `99.37% <0.00%> (-0.32%)` | :arrow_down: |
| [pennylane/devices/default\_qubit.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3790?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdC5weQ==) | `99.69% <0.00%> (-0.01%)` | :arrow_down: |
| [pennylane/utils.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3790?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3V0aWxzLnB5) | `100.00% <0.00%> (ø)` | |
| [pennylane/drawer/draw.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3790?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RyYXdlci9kcmF3LnB5) | `100.00% <0.00%> (ø)` | |
| [pennylane/pulse/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3790?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3B1bHNlL19faW5pdF9fLnB5) | `100.00% <0.00%> (ø)` | |
| [pennylane/qchem/molecule.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3790?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3FjaGVtL21vbGVjdWxlLnB5) | `100.00% <0.00%> (ø)` | |
| ... and [29 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3790?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | |

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

trbromley: [SC-31124]

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 6745540b9..77184f803 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -8,6 +8,10 @@
   state vector and a number of shots.
   [(#3720)](https://github.com/PennyLaneAI/pennylane/pull/3720)
 
+* Adjoint differentiation support for the new qubit state-vector device has been added via
+  `adjoint_jacobian` in `devices/qubit`.
+  [(#3790)](https://github.com/PennyLaneAI/pennylane/pull/3790)
+
 * Added the needed functions and classes to simulate an ensemble of Rydberg atoms:
   * A new internal `RydbergHamiltonian` class is added, which contains the Hamiltonian of an ensemble of
     Rydberg atoms.
@@ -168,7 +172,7 @@
 <h3>Documentation</h3>
 
 * A typo was corrected in the documentation for introduction to `inspecting_circuits` and `chemistry`.
-[(#3844)](https://github.com/PennyLaneAI/pennylane/pull/3844)
+  [(#3844)](https://github.com/PennyLaneAI/pennylane/pull/3844)
 
 <h3>Bug fixes</h3>
 
diff --git a/pennylane/devices/qubit/__init__.py b/pennylane/devices/qubit/__init__.py
index 7e93f236d..98418aef7 100644
--- a/pennylane/devices/qubit/__init__.py
+++ b/pennylane/devices/qubit/__init__.py
@@ -29,10 +29,10 @@ at your own discretion.
     simulate
 """
 
-from .initialize_state import create_initial_state
 from .apply_operation import apply_operation
+from .adjoint_jacobian import adjoint_jacobian
+from .initialize_state import create_initial_state
 from .measure import measure
-from .sampling import sample_state
-
 from .preprocess import preprocess
+from .sampling import sample_state
 from .simulate import simulate
diff --git a/pennylane/devices/qubit/adjoint_jacobian.py b/pennylane/devices/qubit/adjoint_jacobian.py
new file mode 100644
index 000000000..62adbd868
--- /dev/null
+++ b/pennylane/devices/qubit/adjoint_jacobian.py
@@ -0,0 +1,110 @@
+# Copyright 2018-2023 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Functions to apply adjoint jacobian differentiation"""
+
+import numpy as np
+
+import pennylane as qml
+
+from pennylane.operation import operation_derivative
+from pennylane.tape import QuantumTape
+
+from .apply_operation import apply_operation
+from .initialize_state import create_initial_state
+
+# pylint: disable=protected-access, too-many-branches
+
+
+def _dot_product_real(bra, ket, num_wires):
+    """Helper for calculating the inner product for adjoint differentiation."""
+    # broadcasted inner product not summing over first dimension of the bra tensor
+    sum_axes = tuple(range(1, num_wires + 1))
+    return qml.math.real(qml.math.sum(qml.math.conj(bra) * ket, axis=sum_axes))
+
+
+def adjoint_jacobian(tape: QuantumTape):  # pylint: disable=too-many-statements
+    """Implements the adjoint method outlined in
+    `Jones and Gacon <https://arxiv.org/abs/2009.02823>`__ to differentiate an input tape.
+
+    After a forward pass, the circuit is reversed by iteratively applying adjoint
+    gates to scan backwards through the circuit.
+
+    .. note::
+        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
+
+    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
+    if set(tape.wires) != set(range(tape.num_wires)):
+        wire_map = {w: i for i, w in enumerate(tape.wires)}
+        tape = qml.map_wires(tape, wire_map)
+
+    # Initialization of state
+    prep_operation = None if len(tape._prep) == 0 else tape._prep[0]
+    ket = create_initial_state(
+        wires=tape.wires, prep_operation=prep_operation
+    )  #  ket(0) if prep_operation is None, else
+    for op in tape._ops:
+        ket = apply_operation(op, ket)
+
+    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)
+
+    jac = np.zeros((len(tape.observables), len(tape.trainable_params)))
+
+    param_number = len(tape.get_parameters(trainable_only=False, operations_only=True)) - 1
+    trainable_param_number = len(tape.trainable_params) - 1
+    for op in reversed(tape._ops):
+        adj_op = qml.adjoint(op)
+        ket = apply_operation(adj_op, ket)
+
+        if op.grad_method is not 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)
+                )
+
+                trainable_param_number -= 1
+            param_number -= 1
+
+        for kk in range(n_obs):
+            bras[kk, ...] = apply_operation(adj_op, bras[kk, ...])
+
+    # Post-process the Jacobian matrix for the new return
+    jac = np.squeeze(jac)
+
+    if jac.ndim == 0:
+        return np.array(jac)
+
+    if jac.ndim == 1:
+        return tuple(np.array(j) for j in jac)
+
+    # must be 2-dimensional
+    return tuple(tuple(np.array(j_) for j_ in j) for j in jac)
diff --git a/pennylane/devices/qubit/preprocess.py b/pennylane/devices/qubit/preprocess.py
index e20c9ee91..942be22c1 100644
--- a/pennylane/devices/qubit/preprocess.py
+++ b/pennylane/devices/qubit/preprocess.py
@@ -16,11 +16,12 @@
 that they are supported for execution by a device."""
 # pylint: disable=protected-access
 from typing import Generator, Callable, Tuple
+import warnings
 
 import pennylane as qml
 
 from pennylane.operation import Tensor
-from pennylane.measurements import MidMeasureMP, StateMeasurement
+from pennylane.measurements import MidMeasureMP, StateMeasurement, ExpectationMP
 from pennylane import DeviceError
 
 from ..experimental import ExecutionConfig, DefaultExecutionConfig
@@ -77,6 +78,73 @@ def _operator_decomposition_gen(
 #######################
 
 
+def validate_and_expand_adjoint(
+    circuit: qml.tape.QuantumTape,
+) -> qml.tape.QuantumTape:  # pylint: disable=protected-access
+    """Function for validating that the operations and observables present in the input circuit
+    are valid for adjoint differentiation.
+
+    Args:
+        circuit(.QuantumTape): the tape to validate
+
+    Returns:
+        .QuantumTape: The expanded tape, such that it is supported by adjoint differentiation
+
+    Raises:
+        QuantumFunctionError: if the circuit is invalid for adjoint differentiation
+    """
+    # Check validity of measurements
+    measurements = []
+    for m in circuit.measurements:
+        if not isinstance(m, ExpectationMP):
+            raise DeviceError(
+                "Adjoint differentiation method does not support "
+                f"measurement {m.__class__.__name__}."
+            )
+
+        if not m.obs.has_matrix:
+            raise DeviceError(
+                f"Adjoint differentiation method does not support observable {m.obs.name}."
+            )
+
+        measurements.append(m)
+
+    expanded_ops = []
+    for op in circuit._ops:
+        if op.num_params > 1:
+            if not isinstance(op, qml.Rot):
+                raise DeviceError(
+                    f"The {op} operation is not supported using "
+                    'the "adjoint" differentiation method.'
+                )
+            ops = op.decomposition()
+            expanded_ops.extend(ops)
+        elif not isinstance(op, qml.operation.StatePrep):
+            expanded_ops.append(op)
+
+    prep = circuit._prep[:1]
+
+    trainable_params = []
+    for k in circuit.trainable_params:
+        if hasattr(circuit._par_info[k]["op"], "return_type"):
+            warnings.warn(
+                "Differentiating with respect to the input parameters of "
+                f"{circuit._par_info[k]['op'].name} is not supported with the "
+                "adjoint differentiation method. Gradients are computed "
+                "only with regards to the trainable parameters of the circuit.\n\n Mark "
+                "the parameters of the measured observables as non-trainable "
+                "to silence this warning.",
+                UserWarning,
+            )
+        else:
+            trainable_params.append(k)
+
+    expanded_tape = qml.tape.QuantumScript(expanded_ops, measurements, prep)
+    expanded_tape.trainable_params = trainable_params
+
+    return expanded_tape
+
+
 def expand_fn(circuit: qml.tape.QuantumScript) -> qml.tape.QuantumScript:
     """Method for expanding or decomposing an input circuit.
 
@@ -126,7 +194,9 @@ def expand_fn(circuit: qml.tape.QuantumScript) -> qml.tape.QuantumScript:
     return circuit
 
 
-def batch_transform(circuit: qml.tape.QuantumScript) -> (Tuple[qml.tape.QuantumScript], Callable):
+def batch_transform(
+    circuit: qml.tape.QuantumScript,
+) -> Tuple[Tuple[qml.tape.QuantumScript], Callable]:
     """Apply a differentiable batch transform for preprocessing a circuit
     prior to execution.
 
@@ -169,7 +239,7 @@ def batch_transform(circuit: qml.tape.QuantumScript) -> (Tuple[qml.tape.QuantumS
 def preprocess(
     circuits: Tuple[qml.tape.QuantumScript],
     execution_config: ExecutionConfig = DefaultExecutionConfig,
-) -> (Tuple[qml.tape.QuantumScript], Callable):
+) -> Tuple[Tuple[qml.tape.QuantumScript], Callable]:
     """Preprocess a batch of :class:`~.QuantumTape` objects to make them ready for execution.
 
     This function validates a batch of :class:`~.QuantumTape` objects by transforming and expanding
@@ -190,6 +260,8 @@ def preprocess(
         raise DeviceError("The Python Device does not support finite shots.")
 
     circuits = tuple(expand_fn(c) for c in circuits)
+    if execution_config.gradient_method == "adjoint":
+        circuits = tuple(validate_and_expand_adjoint(c) for c in circuits)
 
     circuits, batch_fn = qml.transforms.map_batch_transform(batch_transform, circuits)
 


Test patch:
diff --git a/tests/devices/qubit/test_adjoint_jacobian.py b/tests/devices/qubit/test_adjoint_jacobian.py
new file mode 100644
index 000000000..c71ee4a02
--- /dev/null
+++ b/tests/devices/qubit/test_adjoint_jacobian.py
@@ -0,0 +1,284 @@
+# Copyright 2018-2023 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Unit and integration tests for the adjoint_jacobian function for DefaultQubit2"""
+import pytest
+import pennylane as qml
+from pennylane.devices.qubit import adjoint_jacobian
+from pennylane.tape import QuantumScript
+import pennylane.numpy as np
+from pennylane.devices.qubit.preprocess import validate_and_expand_adjoint
+
+
+class TestAdjointJacobian:
+    """Tests for adjoint_jacobian"""
+
+    def test_custom_wire_labels(self, tol):
+        """Test that adjoint_jacbonian works as expected when custom wire labels are used."""
+        qs = QuantumScript(
+            [qml.RX(0.123, wires="a"), qml.RY(0.456, wires="b")], [qml.expval(qml.PauliX("a"))]
+        )
+        qs.trainable_params = {0, 1}
+
+        qs_valid = validate_and_expand_adjoint(qs)
+
+        calculated_val = adjoint_jacobian(qs_valid)
+
+        tapes, fn = qml.gradients.finite_diff(qs)
+        results = tuple(qml.devices.qubit.simulate(t) for t in tapes)
+        numeric_val = fn(results)
+        assert np.allclose(calculated_val, numeric_val, atol=tol, rtol=0)
+
+    @pytest.mark.autograd
+    @pytest.mark.parametrize("theta", np.linspace(-2 * np.pi, 2 * np.pi, 7))
+    @pytest.mark.parametrize("G", [qml.RX, qml.RY, qml.RZ])
+    def test_pauli_rotation_gradient(self, G, theta, tol):
+        """Tests that the automatic gradients of Pauli rotations are correct."""
+
+        prep_op = qml.QubitStateVector(
+            np.array([1.0, -1.0], requires_grad=False) / np.sqrt(2), wires=0
+        )
+        qs = QuantumScript(
+            ops=[G(theta, wires=[0])], measurements=[qml.expval(qml.PauliZ(0))], prep=[prep_op]
+        )
+
+        qs.trainable_params = {1}
+        qs_valid = validate_and_expand_adjoint(qs)
+
+        calculated_val = adjoint_jacobian(qs_valid)
+
+        # compare to finite differences
+        tapes, fn = qml.gradients.finite_diff(qs)
+        results = tuple(qml.devices.qubit.simulate(t) for t in tapes)
+        numeric_val = fn(results)
+        assert np.allclose(calculated_val, numeric_val, atol=tol, rtol=0)
+        assert isinstance(calculated_val, np.ndarray)
+
+    @pytest.mark.autograd
+    @pytest.mark.parametrize("theta", np.linspace(-2 * np.pi, 2 * np.pi, 7))
+    def test_Rot_gradient(self, theta, tol):
+        """Tests that the device gradient of an arbitrary Euler-angle-parameterized gate is
+        correct."""
+        params = np.array([theta, theta**3, np.sqrt(2) * theta])
+        prep_op = qml.QubitStateVector(
+            np.array([1.0, -1.0], requires_grad=False) / np.sqrt(2), wires=0
+        )
+
+        qs = QuantumScript(
+            ops=[qml.Rot(*params, wires=[0])],
+            measurements=[qml.expval(qml.PauliZ(0))],
+            prep=[prep_op],
+        )
+
+        qs.trainable_params = {1, 2, 3}
+        qs_valid = validate_and_expand_adjoint(qs)
+
+        calculated_val = adjoint_jacobian(qs_valid)
+
+        # compare to finite differences
+        tapes, fn = qml.gradients.finite_diff(qs)
+        results = tuple(qml.devices.qubit.simulate(t) for t in tapes)
+        numeric_val = fn(results)
+        assert np.allclose(calculated_val, numeric_val, atol=tol, rtol=0)
+        assert isinstance(calculated_val, tuple)
+        assert all(isinstance(val, np.ndarray) for val in calculated_val)
+
+    @pytest.mark.autograd
+    @pytest.mark.parametrize("obs", [qml.PauliY])
+    @pytest.mark.parametrize(
+        "op", [qml.RX(0.4, wires=0), qml.CRZ(1.0, wires=[0, 1]), qml.Rot(0.2, -0.1, 0.2, wires=0)]
+    )
+    def test_gradients(self, op, obs, tol):
+        """Tests that the gradients of circuits match between the finite difference and device
+        methods."""
+
+        ops = [
+            qml.Hadamard(wires=0),
+            qml.RX(0.543, wires=0),
+            qml.CNOT(wires=[0, 1]),
+            op,
+            qml.Rot(1.3, -2.3, 0.5, wires=[0]),
+            qml.RZ(-0.5, wires=0),
+            qml.adjoint(qml.RY(0.5, wires=1)),
+            qml.CNOT(wires=[0, 1]),
+        ]
+        measurements = [qml.expval(obs(wires=0)), qml.expval(qml.PauliZ(wires=1))]
+
+        qs = QuantumScript(ops, measurements)
+        qs.trainable_params = set(range(1, 1 + op.num_params))
+        qs_valid = validate_and_expand_adjoint(qs)
+
+        tapes, fn = qml.gradients.finite_diff(qs)
+        results = tuple(qml.devices.qubit.simulate(t) for t in tapes)
+        numeric_val = fn(results)
+        grad_D = adjoint_jacobian(qs_valid)
+
+        grad_F = np.squeeze(numeric_val)
+
+        assert np.allclose(grad_D, grad_F, atol=tol, rtol=0)
+        assert isinstance(grad_D, tuple)
+
+    def test_multiple_rx_gradient(self, tol):
+        """Tests that the gradient of multiple RX gates in a circuit yields the correct result."""
+        params = np.array([np.pi, np.pi / 2, np.pi / 3])
+
+        qs = QuantumScript(
+            [qml.RX(params[0], wires=0), qml.RX(params[1], wires=1), qml.RX(params[2], wires=2)],
+            [qml.expval(qml.PauliZ(idx)) for idx in range(3)],
+        )
+        qs_valid = validate_and_expand_adjoint(qs)
+
+        # circuit jacobians
+        jacobian = adjoint_jacobian(qs_valid)
+        expected_jacobian = -np.diag(np.sin(params))
+        assert np.allclose(jacobian, expected_jacobian, atol=tol, rtol=0)
+        assert isinstance(jacobian, tuple)
+        assert all(isinstance(j, tuple) for j in jacobian)
+
+    def test_custom_op_gradient(self, tol):
+        """Tests that the gradient of a custom operation that only provides a
+        matrix and a generator works."""
+
+        class MyOp(qml.operation.Operation):
+            """Custom operation that only defines a generator and a matrix representation."""
+
+            num_wires = 1
+
+            def generator(self):
+                """Generator of MyOp, just a multiple of Pauli X."""
+                return qml.PauliX(self.wires) * 1.2
+
+            @staticmethod
+            def compute_matrix(angle):  # pylint: disable=arguments-differ
+                """Matrix representation of MyOp, just the same as a reparametrized RX."""
+                return qml.RX.compute_matrix(-2.4 * angle)
+
+        params = np.array([np.pi, np.pi / 2, np.pi / 3])
+
+        qs = QuantumScript(
+            [MyOp(p, w) for p, w in zip(params, [0, 1, 2])],
+            [qml.expval(qml.PauliZ(idx)) for idx in range(3)],
+        )
+        qs_valid = validate_and_expand_adjoint(qs)
+
+        # circuit jacobians
+        jacobian = adjoint_jacobian(qs_valid)
+        expected_jacobian = 2.4 * np.diag(np.sin(-2.4 * params))
+        assert np.allclose(jacobian, expected_jacobian, atol=tol, rtol=0)
+        assert isinstance(jacobian, tuple)
+        assert all(isinstance(j, tuple) for j in jacobian)
+
+    @pytest.mark.autograd
+    def test_gradient_gate_with_multiple_parameters(self, tol):
+        """Tests that gates with multiple free parameters yield correct gradients."""
+        x, y, z = [0.5, 0.3, -0.7]
+
+        qs = QuantumScript(
+            [qml.RX(0.4, wires=[0]), qml.Rot(x, y, z, wires=[0]), qml.RY(-0.2, wires=[0])],
+            [qml.expval(qml.PauliZ(0))],
+        )
+
+        qs.trainable_params = {1, 2, 3}
+        qs_valid = validate_and_expand_adjoint(qs)
+
+        grad_D = adjoint_jacobian(qs_valid)
+        tapes, fn = qml.gradients.finite_diff(qs)
+        results = tuple(qml.devices.qubit.simulate(t) for t in tapes)
+        grad_F = fn(results)
+        grad_F = np.squeeze(grad_F)
+
+        # gradient has the correct shape and every element is nonzero
+        assert len(grad_D) == 3
+        assert np.count_nonzero(grad_D) == 3
+        # the different methods agree
+        assert np.allclose(grad_D, grad_F, atol=tol, rtol=0)
+
+    @pytest.mark.parametrize(
+        "prep_op", [qml.BasisState([1], wires=0), qml.QubitStateVector([0, 1], wires=0)]
+    )
+    def test_state_prep(self, prep_op, tol):
+        """Tests provides correct answer when provided state preparation operation."""
+        x, y, z = [0.5, 0.3, -0.7]
+
+        qs = QuantumScript(
+            [qml.RX(0.4, wires=[0]), qml.Rot(x, y, z, wires=[0]), qml.RY(-0.2, wires=[0])],
+            [qml.expval(qml.PauliZ(0))],
+            [prep_op],
+        )
+
+        qs.trainable_params = {1, 2, 3}
+        qs_valid = validate_and_expand_adjoint(qs)
+
+        grad_D = adjoint_jacobian(qs_valid)
+        tapes, fn = qml.gradients.finite_diff(qs)
+        results = tuple(qml.devices.qubit.simulate(t) for t in tapes)
+        grad_F = fn(results)
+
+        assert np.allclose(grad_D, grad_F, atol=tol, rtol=0)
+
+    def test_gradient_of_tape_with_hermitian(self, tol):
+        """Test that computing the gradient of a tape that obtains the
+        expectation value of a Hermitian operator works correctly."""
+        a, b, c = [0.5, 0.3, -0.7]
+
+        mx = qml.matrix(qml.PauliX(0) @ qml.PauliY(2))
+        qs = QuantumScript(
+            [
+                qml.RX(a, wires=0),
+                qml.RX(b, wires=1),
+                qml.RX(c, wires=2),
+                qml.CNOT(wires=[0, 1]),
+                qml.CNOT(wires=[1, 2]),
+            ],
+            [qml.expval(qml.Hermitian(mx, wires=[0, 2]))],
+        )
+
+        qs.trainable_params = {0, 1, 2}
+        qs_valid = validate_and_expand_adjoint(qs)
+
+        res = adjoint_jacobian(qs_valid)
+
+        expected = [
+            np.cos(a) * np.sin(b) * np.sin(c),
+            np.cos(b) * np.sin(a) * np.sin(c),
+            np.cos(c) * np.sin(b) * np.sin(a),
+        ]
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_gradient_of_tape_with_tensor(self, tol):
+        """Test that computing the gradient of a tape that obtains the
+        expectation value of a Tensor operator works correctly."""
+        a, b, c = [0.5, 0.3, -0.7]
+
+        qs = QuantumScript(
+            [
+                qml.RX(a, wires=0),
+                qml.RX(b, wires=1),
+                qml.RX(c, wires=2),
+                qml.CNOT(wires=[0, 1]),
+                qml.CNOT(wires=[1, 2]),
+            ],
+            [qml.expval(qml.PauliX(0) @ qml.PauliY(2))],
+        )
+
+        qs.trainable_params = {0, 1, 2}
+        qs_valid = validate_and_expand_adjoint(qs)
+
+        res = adjoint_jacobian(qs_valid)
+
+        expected = [
+            np.cos(a) * np.sin(b) * np.sin(c),
+            np.cos(b) * np.sin(a) * np.sin(c),
+            np.cos(c) * np.sin(b) * np.sin(a),
+        ]
+        assert np.allclose(res, expected, atol=tol, rtol=0)
diff --git a/tests/devices/qubit/test_preprocess.py b/tests/devices/qubit/test_preprocess.py
index 60f17c733..38fbdbb19 100644
--- a/tests/devices/qubit/test_preprocess.py
+++ b/tests/devices/qubit/test_preprocess.py
@@ -16,6 +16,7 @@
 import pytest
 
 import numpy as np
+import pennylane.numpy as pnp
 import pennylane as qml
 from pennylane.operation import Operation
 from pennylane.devices.qubit.preprocess import (
@@ -24,6 +25,7 @@ from pennylane.devices.qubit.preprocess import (
     expand_fn,
     batch_transform,
     preprocess,
+    validate_and_expand_adjoint,
 )
 from pennylane.measurements import MidMeasureMP, MeasurementValue
 from pennylane.tape import QuantumScript
@@ -273,6 +275,109 @@ class TestBatchTransform:
         assert np.array_equal(batch_fn(input), np.array([[1, 2], [3, 4]]))
 
 
+class TestAdjointDiffTapeValidation:
+    """Unit tests for validate_and_expand_adjoint"""
+
+    def test_not_expval(self):
+        """Test if a QuantumFunctionError is raised for a tape with measurements that are not
+        expectation values"""
+
+        measurements = [qml.expval(qml.PauliZ(0)), qml.var(qml.PauliX(3)), qml.sample()]
+        qs = QuantumScript(ops=[], measurements=measurements)
+
+        with pytest.raises(
+            DeviceError,
+            match="Adjoint differentiation method does not support measurement",
+        ):
+            validate_and_expand_adjoint(qs)
+
+    def test_unsupported_op(self):
+        """Test if a QuantumFunctionError is raised for an unsupported operation, i.e.,
+        multi-parameter operations that are not qml.Rot"""
+
+        qs = QuantumScript([qml.U2(0.1, 0.2, wires=[0])], [qml.expval(qml.PauliZ(2))])
+
+        with pytest.raises(
+            DeviceError,
+            match='operation is not supported using the "adjoint" differentiation method',
+        ):
+            validate_and_expand_adjoint(qs)
+
+    @pytest.mark.parametrize(
+        "obs",
+        [
+            qml.Hamiltonian([2, 0.5], [qml.PauliZ(0), qml.PauliY(1)]),
+        ],
+    )
+    def test_unsupported_obs(self, obs):
+        """Test that the correct error is raised if a Hamiltonian or Sum measurement is differentiated"""
+        qs = QuantumScript([qml.RX(0.5, wires=1)], [qml.expval(obs)])
+
+        with pytest.raises(
+            DeviceError,
+            match="Adjoint differentiation method does not support observable",
+        ):
+            validate_and_expand_adjoint(qs)
+
+    def test_trainable_hermitian_warns(self):
+        """Test attempting to compute the gradient of a tape that obtains the
+        expectation value of a Hermitian operator emits a warning if the
+        parameters to Hermitian are trainable."""
+
+        mx = qml.matrix(qml.PauliX(0) @ qml.PauliY(2))
+        qs = QuantumScript([], [qml.expval(qml.Hermitian(mx, wires=[0, 2]))])
+
+        qs.trainable_params = {0}
+
+        with pytest.warns(
+            UserWarning, match="Differentiating with respect to the input parameters of Hermitian"
+        ):
+            _ = validate_and_expand_adjoint(qs)
+
+    @pytest.mark.parametrize("G", [qml.RX, qml.RY, qml.RZ])
+    def test_valid_tape_no_expand(self, G):
+        """Test that a tape that is valid doesn't raise errors and is not expanded"""
+        prep_op = qml.QubitStateVector(
+            pnp.array([1.0, -1.0], requires_grad=False) / np.sqrt(2), wires=0
+        )
+        qs = QuantumScript(
+            ops=[G(np.pi, wires=[0])], measurements=[qml.expval(qml.PauliZ(0))], prep=[prep_op]
+        )
+
+        qs.trainable_params = {1}
+        qs_valid = validate_and_expand_adjoint(qs)
+
+        assert all(qml.equal(o1, o2) for o1, o2 in zip(qs.operations, qs_valid.operations))
+        assert all(qml.equal(o1, o2) for o1, o2 in zip(qs.measurements, qs_valid.measurements))
+        assert qs.trainable_params == qs_valid.trainable_params
+
+    def test_valid_tape_with_expansion(self):
+        """Test that a tape that is valid with operations that need to be expanded doesn't raise errors
+        and is expanded"""
+        prep_op = qml.QubitStateVector(
+            pnp.array([1.0, -1.0], requires_grad=False) / np.sqrt(2), wires=0
+        )
+        qs = QuantumScript(
+            ops=[qml.Rot(0.1, 0.2, 0.3, wires=[0])],
+            measurements=[qml.expval(qml.PauliZ(0))],
+            prep=[prep_op],
+        )
+
+        qs.trainable_params = {1, 2, 3}
+        qs_valid = validate_and_expand_adjoint(qs)
+
+        expected_ops = [
+            prep_op,
+            qml.RZ(0.1, wires=[0]),
+            qml.RY(0.2, wires=[0]),
+            qml.RZ(0.3, wires=[0]),
+        ]
+
+        assert all(qml.equal(o1, o2) for o1, o2 in zip(qs_valid.operations, expected_ops))
+        assert all(qml.equal(o1, o2) for o1, o2 in zip(qs.measurements, qs_valid.measurements))
+        assert qs.trainable_params == qs_valid.trainable_params
+
+
 class TestPreprocess:
     """Unit tests for ``qml.devices.qubit.preprocess``."""
 
@@ -377,3 +482,57 @@ class TestPreprocess:
 
         with pytest.raises(qml.DeviceError, match="Operator NoMatNoDecompOp"):
             _ = preprocess(tapes)
+
+    @pytest.mark.parametrize(
+        "ops, measurement, message",
+        [
+            (
+                [qml.RX(0.1, wires=0)],
+                [qml.probs(wires=[0, 1, 2])],
+                "does not support measurement ProbabilityMP",
+            ),
+            (
+                [qml.RX(0.1, wires=0)],
+                [qml.expval(qml.Hamiltonian([1], [qml.PauliZ(0)]))],
+                "does not support observable Hamiltonian",
+            ),
+            (
+                [qml.U2(0.1, 0.2, wires=0)],
+                [qml.expval(qml.PauliZ(0))],
+                'operation is not supported using the "adjoint" differentiation method',
+            ),
+        ],
+    )
+    def test_preprocess_invalid_tape_adjoint(self, ops, measurement, message):
+        """Test that preprocessing fails if adjoint differentiation is requested and an
+        invalid tape is used"""
+        qs = QuantumScript(ops, measurement)
+        execution_config = qml.devices.experimental.ExecutionConfig(gradient_method="adjoint")
+
+        with pytest.raises(DeviceError, match=message):
+            _ = preprocess([qs], execution_config)
+
+    def test_preprocess_tape_for_adjoint(self):
+        """Test that a tape is expanded correctly if adjoint differentiation is requested"""
+        qs = QuantumScript(
+            [qml.Rot(0.1, 0.2, 0.3, wires=0), qml.CNOT([0, 1])], [qml.expval(qml.PauliZ(1))]
+        )
+        execution_config = qml.devices.experimental.ExecutionConfig(gradient_method="adjoint")
+
+        expanded_tapes, _ = preprocess([qs], execution_config)
+
+        assert len(expanded_tapes) == 1
+        expanded_qs = expanded_tapes[0]
+
+        expected_qs = QuantumScript(
+            [qml.RZ(0.1, wires=0), qml.RY(0.2, wires=0), qml.RZ(0.3, wires=0), qml.CNOT([0, 1])],
+            [qml.expval(qml.PauliZ(1))],
+        )
+
+        assert all(
+            qml.equal(o1, o2) for o1, o2 in zip(expanded_qs.operations, expected_qs.operations)
+        )
+        assert all(
+            qml.equal(o1, o2) for o1, o2 in zip(expanded_qs.measurements, expected_qs.measurements)
+        )
+        assert expanded_qs.trainable_params == expected_qs.trainable_params
diff --git a/tests/legacy/qnn/test_keras_old.py b/tests/legacy/qnn/test_keras_old.py
index d4b845ec5..b551310f4 100644
--- a/tests/legacy/qnn/test_keras_old.py
+++ b/tests/legacy/qnn/test_keras_old.py
@@ -14,6 +14,7 @@
 """
 Tests for the pennylane.qnn.keras module.
 """
+# pylint: disable=unnecessary-dunder-call
 import numpy as np
 import pytest
 
diff --git a/tests/legacy/qnn/test_qnn_torch_old.py b/tests/legacy/qnn/test_qnn_torch_old.py
index 460ab65ce..094fcf7e0 100644
--- a/tests/legacy/qnn/test_qnn_torch_old.py
+++ b/tests/legacy/qnn/test_qnn_torch_old.py
@@ -14,6 +14,7 @@
 """
 Tests for the pennylane.qnn.torch module.
 """
+# pylint: disable=unnecessary-dunder-call
 from unittest import mock
 import math
 import numpy as np