Repository: PennyLaneAI/pennylane
Pull request #5829 (2024-06-10 17:36:01)

Add diagonalize_measurements transform
**Context:**
This transform diagonalizes measurements on a tape. By default, it diagonalizes all measurements, but generally this can be done more efficiently using the `pauli_rep` of the observables (next PR). This transform is primarily useful for the added functionality of only diagonalizing _some subset_ of {X, Y, Z and Hadamard}, for backends that support only some of these. This is based on requirements established for Catalyst backends.

**Description of the Change:**
The `diagonalize_tape_measurements` transform is added. If the backend supports some non-measurement basis observables, we can diagonalize only a subset by passing `supported_base_obs`:

```
>>> measurements = [qml.expval(qml.X(0) @ qml.Y(1)), qml.var(qml.X(0))]
>>> tape = qml.tape.QuantumScript([], measurements=measurements)
>>> tapes, _ = diagonalize_tape_measurements(tape, supported_base_obs =["PauliY", "PauliZ"])
>>> tapes[0].measurements
[expval(Z(0) @ Y(1)), var(Z(0))]
```

The transform tracks which observables and wires are already encountered, and avoids applying diagonalizing gates more than once for the same wire:

```
>>> tapes[0].operations
[Hadamard(wires=[0])]
```

**Benefits:**
We can diagonalize a subset of observables on a tape.

**Possible Drawbacks:**
- If we are generally diagonalizing (the most common use for diagonalization), for larger observables performance doesn't scale well compared to a pauli_rep based approach. We will also implement this (followup PR).

- This converts all observables to Z instead of to eigvals. It would be nice to have the option to chose either one, but that is out of scope and would need to be a future improvement.




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.
trbromley: @lillian542 is there a story or issue that can be associated with this?
lillian542: [sc-60179]
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5829?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: Patch coverage is `0%` with `81 lines` in your changes missing coverage. Please review.
> Project coverage is 99.46%. Comparing base [(`b7675ca`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/b7675ca42f2e237d396788b287de256254a71958?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`dd85c41`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/dd85c411088baed2adc8b6230862eda37f0dbdd7?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 7 commits behind head on master.

| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5829?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Patch % | Lines |
|---|---|---|
| [pennylane/transforms/diagonalize\_measurements.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5829?src=pr&el=tree&filepath=pennylane%2Ftransforms%2Fdiagonalize_measurements.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvZGlhZ29uYWxpemVfbWVhc3VyZW1lbnRzLnB5) | 0.00% | [81 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5829?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |

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


```diff
@@            Coverage Diff             @@
##           master    #5829      +/-   ##
==========================================
- Coverage   99.66%   99.46%   -0.20%     
==========================================
  Files         428      431       +3     
  Lines       41182    41250      +68     
==========================================
- Hits        41042    41028      -14     
- Misses        140      222      +82     
```

</details>

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

astralcai: > What happens if someone does something like `diagonalize_tape_measurements(tape, supported_base_obs = {"PauliY"})`, and it doesn't include `PauilZ`? Do we always assume `PauliZ` is supported?

Is there a device out there that does not use Pauli Z as its computational basis? 
lillian542: >  What happens if someone does something like diagonalize_tape_measurements(tape, supported_base_obs = {"PauliY"}), and it doesn't include PauilZ? Do we always assume PauliZ is supported?

That's a good point, I didn't consider that people might omit that. I guess we could either assume it's supported, or raise an error. The same goes for `Identity`... I'm kind of in favour of just adding them on there and always supporting them, regardless of user input. That's what I've done for now, lmk what you think.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 4dd6effac..6dc09c6f5 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -65,6 +65,11 @@
   users to have more control over the error mitigation protocol without needing to add further dependencies.
   [(#5972)](https://github.com/PennyLaneAI/pennylane/pull/5972)
 
+* The `diagonalize_measurements` transform is added. This transform converts measurements
+  to the Z basis by applying the relevant diagonalizing gates. It can be set to diagonalize only 
+  a subset of the base observables `{X, Y, Z, Hadamard}`.
+  [(#5829)](https://github.com/PennyLaneAI/pennylane/pull/5829)
+
 * The `qml.PrepSelPrep` template is added. The template implements a block-encoding of a linear
   combination of unitaries.
   [(#5756)](https://github.com/PennyLaneAI/pennylane/pull/5756)
diff --git a/pennylane/transforms/__init__.py b/pennylane/transforms/__init__.py
index b232176ec..e6870d8d7 100644
--- a/pennylane/transforms/__init__.py
+++ b/pennylane/transforms/__init__.py
@@ -110,6 +110,7 @@ preprocessing, getting information from a circuit, and more.
     ~transforms.insert
     ~transforms.add_noise
     ~defer_measurements
+    ~transforms.diagonalize_measurements
     ~transforms.split_non_commuting
     ~transforms.split_to_single_terms
     ~transforms.broadcast_expand
@@ -316,6 +317,7 @@ from .add_noise import add_noise
 
 from .decompositions import clifford_t_decomposition
 from .defer_measurements import defer_measurements
+from .diagonalize_measurements import diagonalize_measurements
 from .dynamic_one_shot import dynamic_one_shot, is_mcm
 from .sign_expand import sign_expand
 from .hamiltonian_expand import hamiltonian_expand, sum_expand
diff --git a/pennylane/transforms/diagonalize_measurements.py b/pennylane/transforms/diagonalize_measurements.py
new file mode 100644
index 000000000..07a415c8a
--- /dev/null
+++ b/pennylane/transforms/diagonalize_measurements.py
@@ -0,0 +1,343 @@
+"""Transform to diagonalize measurements on a tape, assuming all measurements are commuting."""
+
+from copy import copy
+from functools import singledispatch
+
+import pennylane as qml
+from pennylane.operation import Tensor
+from pennylane.ops import CompositeOp, LinearCombination, SymbolicOp
+from pennylane.transforms.core import transform
+
+# pylint: disable=protected-access
+
+_default_supported_obs = (qml.Z, qml.Identity)
+
+
+def null_postprocessing(results):
+    """A postprocessing function returned by a transform that only converts the batch of results
+    into a result for a single ``QuantumTape``.
+    """
+    return results[0]
+
+
+@transform
+def diagonalize_measurements(tape, supported_base_obs=_default_supported_obs):
+    """Diagonalize a set of measurements into the standard basis. Raises an error if the
+    measurements do not commute.
+
+    See the usage details for more information on which measurements are supported.
+
+    Args:
+        tape (QNode or QuantumScript or Callable): The quantum circuit to modify the measurements of.
+        supported_base_obs (Optional, Iterable(Observable)): A list of supported base observable classes.
+            Allowed observables are X, Y, Z, Hadamarda and Identity. Z and Identity are always treated
+            as supported, regardless of input. If no list is provided, the transform will diagonalize
+            everything into the Z basis. If a list is provided, only unsupported observables will be
+            diagonalized to the Z basis.
+
+    Returns:
+        qnode (QNode) or tuple[List[QuantumScript], function]: The transformed circuit as described in :func:`qml.transform <pennylane.transform>`.
+
+    .. note::
+        This transform will raise an error if it encounters non-commuting terms. To avoid non-commuting terms in
+        circuit measurements, the :func:`split_non_commuting <pennylane.transforms.split_non_commuting>` transform
+        can be applied.
+
+
+    **Examples:**
+
+    This transform allows us to transform QNode measurements into the measurement basis by adding
+    the relevant diagonalizing gates to the end of the tape operations.
+
+    .. code-block:: python3
+
+        from pennylane.transforms import diagonalize_measurements
+
+        dev = qml.device("default.qubit")
+
+        @diagonalize_measurements
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RY(x[0], wires=0)
+            qml.RX(x[1], wires=1)
+            return qml.expval(qml.X(0) @ qml.Z(1)), qml.var(0.5 * qml.Y(2) + qml.X(0))
+
+    Applying the transform appends the relevant gates to the end of the circuit to allow
+    measurements to be in the Z basis, so the original circuit
+
+    >>> print(qml.draw(circuit, level=0)([np.pi/4, np.pi/4]))
+    0: ──RY(0.79)─┤ ╭<X@Z> ╭Var[(0.50*Y)+X]
+    1: ──RX(0.79)─┤ ╰<X@Z> │
+    2: ───────────┤        ╰Var[(0.50*Y)+X]
+
+    becomes
+
+    >>> print(qml.draw(circuit)([np.pi/4, np.pi/4]))
+    0: ──RY(0.79)──H────┤ ╭<Z@Z> ╭Var[(0.50*Z)+Z]
+    1: ──RX(0.79)───────┤ ╰<Z@Z> │
+    2: ──Z─────────S──H─┤        ╰Var[(0.50*Z)+Z]
+
+    >>> circuit([np.pi/4, np.pi/4])
+    (tensor(0.5, requires_grad=True), tensor(0.75, requires_grad=True))
+
+    .. details::
+        :title: Usage Details
+
+        The transform diagonalizes observables from the local Pauli basis only, i.e. it diagonalizes X, Y, Z,
+        and Hadamard.
+
+        The transform can also diagonalize only a subset of these operators. By default, the only
+        supported base observable is Z. What if a backend device can handle
+        X, Y and Z, but doesn't provide support for Hadamard? We can set this by passing
+        ``supported_base_obs`` to the transform. Let's create a tape with some measurements:
+
+        .. code-block:: python3
+
+            measurements = [
+                qml.expval(qml.X(0) + qml.Hadamard(1)),
+                qml.expval(qml.X(0) + 0.2 * qml.Hadamard(1)),
+                qml.var(qml.Y(2) + qml.X(0)),
+            ]
+            tape = qml.tape.QuantumScript(measurements=measurements)
+            tapes, processing_fn = diagonalize_measurements(
+                tape,
+                supported_base_obs=[qml.X, qml.Y, qml.Z]
+            )
+
+        Now ``tapes`` is a tuple containing a single tape with the updated measurements,
+        where only the Hadamard gate has been diagonalized:
+
+        >>> tapes[0].measurements
+        [expval(X(0) + Z(1)), expval(X(0) + 0.2 * Z(1)), var(Y(2) + X(0))]
+    """
+
+    bad_obs_input = [
+        o for o in supported_base_obs if o not in {qml.X, qml.Y, qml.Z, qml.Hadamard, qml.Identity}
+    ]
+
+    if bad_obs_input:
+        raise ValueError(
+            "Supported base observables must be a subset of [X, Y, Z, Hadamard, and Identity] "
+            f"but received {list(bad_obs_input)}"
+        )
+
+    supported_base_obs = set(list(supported_base_obs) + [qml.Z, qml.Identity])
+
+    _visited_obs = (set(), set())  # tracks which observables and wires have been diagonalized
+    diagonalizing_gates = []
+    new_measurements = []
+
+    for m in tape.measurements:
+        if m.obs:
+            gates, new_obs, _visited_obs = _diagonalize_observable(
+                m.obs, _visited_obs, supported_base_obs
+            )
+            diagonalizing_gates.extend(gates)
+
+            meas = copy(m)
+            meas.obs = new_obs
+            new_measurements.append(meas)
+        else:
+            new_measurements.append(m)
+
+    new_operations = tape.operations + diagonalizing_gates
+
+    new_tape = type(tape)(
+        ops=new_operations,
+        measurements=new_measurements,
+        shots=tape.shots,
+        trainable_params=tape.trainable_params,
+    )
+
+    return (new_tape,), null_postprocessing
+
+
+def _check_if_diagonalizing(obs, _visited_obs, switch_basis):
+    """Checks if the observable should be diagonalized based on whether its basis should
+    be switched, and whether the same observable has already been diagonalized.
+
+    Args:
+        obs: the observable to be diagonalized
+        _visited_obs: a tuple containing a list of all observables that have already
+            been encountered, and a list of all wires that have already been encountered
+        switch_basis: whether the observable should be switched for one measuring in
+            the Z-basis on the same wires.
+
+    Returns:
+        diagonalize (bool): whether or not to apply diagonalizing gates for the observable
+        _visited_obs (tuple(Iterables): an up-to-date record of the observables and wires
+            encountered on the tape so far
+    """
+
+    if obs in _visited_obs[0] or isinstance(obs, qml.Identity):
+        # its already been encountered before, and if need be, diagonalized - we will
+        # not be applying any gates or updating _visited_obs
+        # same if it's just an Identity
+        return False, _visited_obs
+
+    # a different observable has been diagonalized on the same wire - error
+    if obs.wires[0] in _visited_obs[1]:
+        raise ValueError(
+            f"Expected measurements on the same wire to commute, but {obs} "
+            f"overlaps with another non-commuting observable on the tape."
+        )
+
+    # we diagonalize if it's an operator we are switching the basis for
+    # we update _visited_obs to indicate that we've encountered that observable regardless
+    _visited_obs[0].add(obs)
+    _visited_obs[1].add(obs.wires[0])
+
+    return switch_basis, _visited_obs
+
+
+def _diagonalize_observable(
+    observable,
+    _visited_obs=None,
+    supported_base_obs=_default_supported_obs,
+):
+    """Takes an observable and changes all unsupported obs to the measurement
+    basis. Applies diagonalizing gates if the observable being switched to the
+    measurement basis hasn't already been diagonalized for a previous observable
+    on the tape.
+
+    Args:
+        observable: the observable to be diagonalized
+        _visited_obs: a tuple containing a list of all observables that have already
+            been encountered, and a list of all wires that have already been encountered
+        supported_base_obs (Optional, Iterable(Str)): A list of names of supported base observables.
+            Allowed names are 'PauliX', 'PauliY', 'PauliZ' and 'Hadamard'. If no list is provided,
+            the function will diagonalize everything.
+
+    Returns:
+        diagonalizing_gates: A list of operations to be applied to diagonalize the observable
+        new_obs: the relevant measurement to perform after applying diagonalzing_gates to get the
+            correct measurement output
+        _visited_obs: an up-to-date record of the observables and wires
+            encountered on the tape so far
+    """
+
+    if _visited_obs is None:
+        _visited_obs = (set(), set())
+
+    if not isinstance(observable, (qml.X, qml.Y, qml.Z, qml.Hadamard, qml.Identity)):
+        return _diagonalize_compound_observable(
+            observable, _visited_obs, supported_base_obs=supported_base_obs
+        )
+
+    # also validates that the wire hasn't already been diagonalized and updated _visited_obs
+    switch_basis = type(observable) not in supported_base_obs
+    print(supported_base_obs)
+    print(type(observable))
+    print(switch_basis)
+    diagonalize, _visited_obs = _check_if_diagonalizing(observable, _visited_obs, switch_basis)
+
+    if isinstance(observable, qml.Z):  # maybe kind of redundant
+        return [], observable, _visited_obs
+
+    new_obs = qml.Z(wires=observable.wires) if switch_basis else observable
+    diagonalizing_gates = observable.diagonalizing_gates() if diagonalize else []
+
+    return diagonalizing_gates, new_obs, _visited_obs
+
+
+def _get_obs_and_gates(obs_list, _visited_obs, supported_base_obs=_default_supported_obs):
+    """Calls _diagonalize_observable on each observable in a list, and returns the full result
+    for the list of observables"""
+
+    new_obs = []
+    diagonalizing_gates = []
+
+    for o in obs_list:
+        gates, obs, _visited_obs = _diagonalize_observable(o, _visited_obs, supported_base_obs)
+        if gates:
+            diagonalizing_gates.extend(gates)
+        new_obs.append(obs)
+
+    return diagonalizing_gates, new_obs, _visited_obs
+
+
+@singledispatch
+def _diagonalize_compound_observable(
+    observable, _visited_obs, supported_base_obs=_default_supported_obs
+):
+    """Takes an observable consisting of multiple other observables, and changes all
+    unsupported obs to the measurement basis. Applies diagonalizing gates if changing
+    the basis of an observable whose diagonalizing gates have not already been applied."""
+
+    raise NotImplementedError(
+        f"Unable to convert observable of type {type(observable)} to the measurement basis"
+    )
+
+
+@_diagonalize_compound_observable.register
+def _diagonalize_symbolic_op(
+    observable: SymbolicOp, _visited_obs, supported_base_obs=_default_supported_obs
+):
+    diagonalizing_gates, new_base, _visited_obs = _diagonalize_observable(
+        observable.base, _visited_obs, supported_base_obs
+    )
+
+    params, hyperparams = observable.parameters, observable.hyperparameters
+    hyperparams = copy(hyperparams)
+    hyperparams["base"] = new_base
+
+    new_observable = observable.__class__(*params, **hyperparams)
+
+    return diagonalizing_gates, new_observable, _visited_obs
+
+
+@_diagonalize_compound_observable.register
+def _diagonalize_tensor(
+    observable: Tensor, _visited_obs, supported_base_obs=_default_supported_obs
+):
+    diagonalizing_gates, new_obs, _visited_obs = _get_obs_and_gates(
+        observable.obs, _visited_obs, supported_base_obs
+    )
+
+    new_observable = Tensor(*new_obs)
+
+    return diagonalizing_gates, new_observable, _visited_obs
+
+
+@_diagonalize_compound_observable.register
+def _diagonalize_hamiltonian(
+    observable: qml.ops.Hamiltonian,
+    _visited_obs,
+    supported_base_obs=_default_supported_obs,
+):
+    diagonalizing_gates, new_ops, _visited_obs = _get_obs_and_gates(
+        observable.ops, _visited_obs, supported_base_obs
+    )
+
+    new_observable = qml.ops.Hamiltonian(observable.coeffs, new_ops)
+
+    return diagonalizing_gates, new_observable, _visited_obs
+
+
+@_diagonalize_compound_observable.register
+def _diagonalize_composite_op(
+    observable: LinearCombination, _visited_obs, supported_base_obs=_default_supported_obs
+):
+
+    coeffs, obs = observable.terms()
+
+    diagonalizing_gates, new_obs, _visited_obs = _get_obs_and_gates(
+        obs, _visited_obs, supported_base_obs
+    )
+
+    new_observable = LinearCombination(coeffs, new_obs)
+
+    return diagonalizing_gates, new_observable, _visited_obs
+
+
+@_diagonalize_compound_observable.register
+def _diagonalize_composite_op(
+    observable: CompositeOp, _visited_obs, supported_base_obs=_default_supported_obs
+):
+    diagonalizing_gates, new_operands, _visited_obs = _get_obs_and_gates(
+        observable.operands, _visited_obs, supported_base_obs
+    )
+
+    new_observable = observable.__class__(*new_operands)
+
+    return diagonalizing_gates, new_observable, _visited_obs


Test patch:
diff --git a/tests/transforms/test_diagonalize_measurements.py b/tests/transforms/test_diagonalize_measurements.py
new file mode 100644
index 000000000..792c2d596
--- /dev/null
+++ b/tests/transforms/test_diagonalize_measurements.py
@@ -0,0 +1,418 @@
+# Copyright 2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Tests for the transform diagonalize_measurements, which diagonalizes unsupported
+observables in measurements on a tape.
+"""
+
+from functools import partial
+
+import numpy as np
+import pytest
+
+import pennylane as qml
+from pennylane import Hadamard, X, Y, Z
+from pennylane.tape import QuantumScript
+from pennylane.transforms.diagonalize_measurements import (
+    _check_if_diagonalizing,
+    _diagonalize_observable,
+    diagonalize_measurements,
+    null_postprocessing,
+)
+
+
+class TestDiagonalizeObservable:
+    """Tests for the _diagonalize_observable method"""
+
+    @pytest.mark.parametrize("obs", [X(1), Y(2), Z(3), Hadamard(0)])
+    def test_diagonalize_observable_single_observable(self, obs):
+        """Test the diagonalize_observables function for X, Y, Z and H"""
+
+        diagonalizing_gates, new_obs, _diag_obs = _diagonalize_observable(obs)
+
+        assert diagonalizing_gates == obs.diagonalizing_gates()
+        assert new_obs == Z(obs.wires)
+        assert _diag_obs == ({obs}, {obs.wires[0]})
+
+    @pytest.mark.parametrize("obs, diagonalize", [(X(1), False), (Y(3), True), (Hadamard(0), True)])
+    def test_supported_base_obs_arg(self, obs, diagonalize):
+        """Test that observables are diagonalized or not depending on whether the observable
+        type is included in supported_base_obs"""
+
+        device_supported_obs = [qml.X, qml.Z]
+
+        diagonalizing_gates, new_obs, _visited_obs = _diagonalize_observable(
+            obs, supported_base_obs=device_supported_obs
+        )
+        assert _visited_obs == ({obs}, {obs.wires[0]})
+
+        if diagonalize:
+            assert diagonalizing_gates == obs.diagonalizing_gates()
+            assert new_obs == Z(obs.wires)
+        else:
+            assert diagonalizing_gates == []
+            assert new_obs == obs
+
+    @pytest.mark.parametrize("obs, apply_gates", [(X(0), True), (Y(2), True), (Y(3), False)])
+    def test_visited_obs_arg(self, obs, apply_gates):
+        """Test that if _visited_obs includes previously encountered observables that overlap
+        with the observable to be diagonalized, this is taken into account"""
+
+        visited_obs = ({Y(3), Z(1)}, {3, 1})
+
+        diagonalizing_gates, new_obs, new_visited_obs = _diagonalize_observable(
+            obs, _visited_obs=visited_obs
+        )
+
+        if apply_gates:
+            assert diagonalizing_gates == obs.diagonalizing_gates()
+            assert new_obs == Z(obs.wires)
+        else:
+            assert diagonalizing_gates == []
+            assert new_obs == Z(obs.wires)
+
+        if obs in visited_obs[0]:
+            assert new_visited_obs == visited_obs
+        else:
+            assert new_visited_obs == ([Y(3), Z(1), obs], [3, 1, obs.wires[0]])
+
+    @pytest.mark.parametrize(
+        "compound_obs, expected_res, base_obs",
+        [
+            (X(0) @ Y(2), Z(0) @ Z(2), [X(0), Y(2)]),  # prod
+            (
+                qml.operation.Tensor(X(0), Y(2)),
+                qml.operation.Tensor(Z(0), Z(2)),
+                [X(0), Y(2)],
+            ),  # tensor
+            (X(1) + Y(2), Z(1) + Z(2), [X(1), Y(2)]),  # sum
+            (2 * X(1), 2 * Z(1), [X(1)]),  # sprod
+            (
+                qml.ops.LinearCombination([2, 3], [X(0), Y(1)]),
+                qml.ops.LinearCombination([2, 3], [Z(0), Z(1)]),
+                [X(0), Y(1)],
+            ),  # hamiltonian
+            (
+                qml.ops.LinearCombination([2, 3], [X(0) @ Y(2), 3 * Y(1)]),
+                qml.ops.LinearCombination([2, 3], [Z(0) @ Z(2), 3 * Z(1)]),
+                [X(0), Y(2), Y(1)],
+            ),  # hamiltonian with composite terms
+            (
+                X(1) + 2 * Y(3) + X(0) @ (3 * Z(4)) + X(0) @ (Y(3) + Z(4)),
+                Z(1) + 2 * Z(3) + Z(0) @ (3 * Z(4)) + Z(0) @ (Z(3) + Z(4)),
+                [X(1), Y(3), X(0), Z(4)],
+            ),  # nested messy sum
+        ],
+    )
+    def test_compound_observables(self, compound_obs, expected_res, base_obs):
+        """Test that _diagonalize_observable works on compound observables"""
+
+        diagonalizing_gates, new_obs, visited_obs = _diagonalize_observable(compound_obs)
+
+        expected_diag_gates = np.concatenate([o.diagonalizing_gates() for o in base_obs])
+
+        assert new_obs == expected_res
+        assert visited_obs == (set(base_obs), {o.wires[0] for o in base_obs})
+        assert diagonalizing_gates == list(expected_diag_gates)
+
+    def test_legacy_hamiltonian(self):
+        """Test that _diagonalize_observable works on legacy Hamiltonians observables"""
+        with pytest.warns():
+            compound_obs = qml.ops.Hamiltonian([2, 3], [Y(0), X(1)])
+            expected_res = qml.ops.Hamiltonian([2, 3], [Z(0), Z(1)])
+            diagonalizing_gates, new_obs, visited_obs = _diagonalize_observable(compound_obs)
+
+        base_obs = [Y(0), X(1)]
+        expected_diag_gates = np.concatenate([o.diagonalizing_gates() for o in base_obs])
+
+        assert new_obs == expected_res
+        assert visited_obs == (set(base_obs), {o.wires[0] for o in base_obs})
+        assert diagonalizing_gates == list(expected_diag_gates)
+
+    @pytest.mark.parametrize(
+        "compound_obs, expected_res, base_obs",
+        [
+            (X(0) @ Y(2), X(0) @ Z(2), [X(0), Y(2)]),  # prod
+            (
+                qml.operation.Tensor(X(0), Y(2)),
+                qml.operation.Tensor(X(0), Z(2)),
+                [X(0), Y(2)],
+            ),  # tensor
+            (X(1) + Y(2), X(1) + Z(2), [X(1), Y(2)]),  # sum
+            (2 * X(1), 2 * X(1), [X(1)]),  # sprod
+            (
+                qml.ops.LinearCombination([2, 3], [X(0), Y(1)]),
+                qml.ops.LinearCombination([2, 3], [X(0), Z(1)]),
+                [X(0), Y(1)],
+            ),  # hamiltonian
+            (
+                qml.ops.LinearCombination([2, 3], [X(0) @ Y(2), 3 * Y(1)]),
+                qml.ops.LinearCombination([2, 3], [X(0) @ Z(2), 3 * Z(1)]),
+                [X(0), Y(2), Y(1)],
+            ),  # hamiltonian with composite terms
+            (
+                X(1) + 2 * Hadamard(3) + X(0) @ (3 * Z(4)) + X(0) @ (Hadamard(3) + Z(4)),
+                X(1) + 2 * Z(3) + X(0) @ (3 * Z(4)) + X(0) @ (Z(3) + Z(4)),
+                [X(1), Hadamard(3), X(0), Z(4)],
+            ),  # nested messy sum
+        ],
+    )
+    def test_compound_observables_supported_base_obs(self, compound_obs, expected_res, base_obs):
+        """Test supported_base_obs argument works as expected for compound observables"""
+
+        device_supported_obs = (qml.X, qml.Z)
+        diagonalizing_gates, new_obs, visited_obs = _diagonalize_observable(
+            compound_obs, supported_base_obs=device_supported_obs
+        )
+
+        diag_obs = [o for o in base_obs if isinstance(o, (Y, Hadamard))]
+
+        expected_diag_gates = (
+            np.concatenate([o.diagonalizing_gates() for o in diag_obs]) if diag_obs else []
+        )
+
+        assert new_obs == expected_res
+        assert visited_obs == (set(base_obs), {o.wires[0] for o in base_obs})
+        assert diagonalizing_gates == list(expected_diag_gates)
+
+    def test_compound_observable_with_duplicate_terms(self):
+        """Test that when a compound observable includes duplicate terms, it only adds
+        the diagonalizing gates once"""
+
+        obs = X(0) + Z(2) + X(0) @ Z(3)
+
+        diagonalizing_gates, new_obs, new_visited_obs = _diagonalize_observable(obs)
+
+        assert new_visited_obs == ({X(0), Z(2), Z(3)}, {0, 2, 3})
+        assert diagonalizing_gates == X(0).diagonalizing_gates()
+        assert new_obs == Z(0) + Z(2) + Z(0) @ Z(3)
+
+    @pytest.mark.parametrize("obs", [X(0) + 2 * qml.Identity(), X(0) + 2 * qml.Identity(wires=2)])
+    def test_with_identity(self, obs):
+        """Test that observables with Identity are supported and Identity remains unchanged"""
+
+        gates, new_obs, visited_obs = _diagonalize_observable(obs)
+
+        assert gates == X(0).diagonalizing_gates()
+        assert new_obs[0] == Z(0)  # X(0) is diagonalized
+        assert new_obs[1] == obs[1]  # Identity is unchanged
+        assert visited_obs == ({X(0)}, {0})
+
+    @pytest.mark.parametrize(
+        "obs", [X(0) + 1.7 * X(2) + X(0) @ Y(2), X(0) + 2.3 * X(2) + X(0) @ Z(2)]
+    )
+    def test_non_commuting_measurements(self, obs):
+        """Test that when a compound observable includes non-commuting observables, it raises
+        an error"""
+
+        obs = X(0) + Z(2) + X(0) @ Y(2)
+
+        with pytest.raises(ValueError, match="Expected measurements on the same wire to commute"):
+            _ = _diagonalize_observable(obs)
+
+    @pytest.mark.parametrize(
+        "obs", [X(0) + 1.7 * X(2) + X(0) @ Y(2), X(0) + 2.3 * X(2) + X(0) @ Z(2)]
+    )
+    def test_non_commuting_measurements_with_supported_obs(self, obs):
+        """Test that when a compound observable includes non-commuting observables, it raises
+        an error, even if some of those observables aren't being diagonalized"""
+
+        device_supported_obs = ["PauliX", "PauliZ"]
+
+        with pytest.raises(ValueError, match="Expected measurements on the same wire to commute"):
+            _ = _diagonalize_observable(obs, supported_base_obs=device_supported_obs)
+
+    def test_diagonalizing_unknown_observable_raises_error(self):
+        """Test that an unknown observable raises an error when diagonalizing"""
+
+        # pylint: disable=too-few-public-methods
+        class MyObs(qml.operation.Observable):
+
+            @property
+            def name(self):
+                return f"MyObservable[{self.wires}]"
+
+        with pytest.raises(NotImplementedError, match="Unable to convert observable"):
+            _ = _diagonalize_observable(MyObs(wires=[2]))
+
+    @pytest.mark.parametrize(
+        "obs, input_visited_obs, switch_basis, expected_res",
+        [
+            (X(0), (set(), set()), True, (True, ({X(0)}, {0}))),
+            (X(0), (set(), set()), False, (False, ({X(0)}, {0}))),
+            (X(0), ({X(0)}, {0}), False, (False, ({X(0)}, {0}))),
+            (X(0), ({X(0)}, {0}), True, (False, ({X(0)}, {0}))),
+        ],
+    )
+    def test_check_if_diagonalising(self, obs, input_visited_obs, switch_basis, expected_res):
+        """Test that _check_if_diagonalizing returns True or False based on whether the
+        observable is supported, and whether its already been diagonalized previously"""
+        diagonalize, output_visited_obs = _check_if_diagonalizing(
+            obs, input_visited_obs, switch_basis
+        )
+
+        assert (diagonalize, output_visited_obs) == expected_res
+
+    @pytest.mark.parametrize(
+        "obs, _visited_obs, raise_error",
+        [(Y(1), ({X(0)}, {0}), False), (Y(1), ({Y(1)}, {1}), False), (Y(1), ({X(1)}, {1}), True)],
+    )
+    def test_check_if_diagonalizing_raises_error(self, obs, _visited_obs, raise_error):
+        """Test that _check_if_diagonalizing raises an error if the observable should be
+        diagonalized, but a different observable on that wire has already been diagonalized"""
+        if raise_error:
+            with pytest.raises(
+                ValueError, match="Expected measurements on the same wire to commute"
+            ):
+                _ = _check_if_diagonalizing(obs, _visited_obs, switch_basis=True)
+            with pytest.raises(
+                ValueError, match="Expected measurements on the same wire to commute"
+            ):
+                _ = _check_if_diagonalizing(obs, _visited_obs, switch_basis=False)
+
+        else:
+            _ = _check_if_diagonalizing(obs, _visited_obs, switch_basis=True)
+            _ = _check_if_diagonalizing(obs, _visited_obs, switch_basis=False)
+
+
+class TestDiagonalizeTapeMeasurements:
+    """Tests the diagonalize_measurements transform"""
+
+    def test_diagonalize_measurements(self):
+        """Test that the diagonalize_measurements transform diagonalizes the measurements on the tape"""
+        measurements = [qml.expval(X(0)), qml.var(X(1) + Y(2))]
+
+        tape = QuantumScript([], measurements=measurements)
+        tapes, fn = diagonalize_measurements(tape)
+        new_tape = tapes[0]
+
+        assert new_tape.measurements == [qml.expval(Z(0)), qml.var(Z(1) + Z(2))]
+        assert (
+            new_tape.operations
+            == X(0).diagonalizing_gates() + X(1).diagonalizing_gates() + Y(2).diagonalizing_gates()
+        )
+
+        assert fn == null_postprocessing
+
+    def test_with_duplicate_measurements(self):
+        """Test that the diagonalize_measurements transform diagonalizes the measurements
+        on the tape correctly when the same observable is used more than once"""
+        measurements = [qml.expval(X(0)), qml.var(X(1) + Y(2)), qml.sample(X(0) @ Y(2))]
+
+        tape = QuantumScript([], measurements=measurements)
+        tapes, fn = diagonalize_measurements(tape)
+        new_tape = tapes[0]
+
+        assert new_tape.measurements == [
+            qml.expval(Z(0)),
+            qml.var(Z(1) + Z(2)),
+            qml.sample(Z(0) @ Z(2)),
+        ]
+        assert (
+            new_tape.operations
+            == X(0).diagonalizing_gates() + X(1).diagonalizing_gates() + Y(2).diagonalizing_gates()
+        )
+
+        assert fn == null_postprocessing
+
+    def test_non_commuting_observables_raise_an_error(self):
+        """Test that the diagonalize_measurements raises an error as expected if the tape contains
+        non-commuting observables"""
+        measurements = [qml.expval(X(0)), qml.var(Z(0) + Y(2))]
+
+        tape = QuantumScript([], measurements=measurements)
+
+        with pytest.raises(ValueError, match="Expected measurements on the same wire to commute"):
+            _ = diagonalize_measurements(tape)
+
+    def test_measurements_with_no_obs(self):
+        """Test that the transform correctly handles tapes where some measurements don't
+        have an observable"""
+
+        measurements = [qml.expval(X(0)), qml.var(X(1) + Y(2)), qml.sample()]
+
+        tape = QuantumScript([], measurements=measurements)
+        tapes, fn = diagonalize_measurements(tape)
+        new_tape = tapes[0]
+
+        assert new_tape.measurements == [qml.expval(Z(0)), qml.var(Z(1) + Z(2)), qml.sample()]
+        assert (
+            new_tape.operations
+            == X(0).diagonalizing_gates() + X(1).diagonalizing_gates() + Y(2).diagonalizing_gates()
+        )
+
+        assert fn == null_postprocessing
+
+    def test_decomposing_subset_of_obs(self):
+        """Test that passing a list of supported obs to the diagonalize_measurements transform
+        diagonalizes only the unsupported base observables"""
+        measurements = [
+            qml.expval(X(0)),
+            qml.var(X(1) + Y(2) @ X(1)),
+            qml.counts(X(0) @ (2.3 * Y(2))),
+        ]
+
+        tape = QuantumScript([], measurements=measurements)
+
+        tapes, fn = diagonalize_measurements(tape, supported_base_obs=[qml.X, qml.Z])
+
+        new_tape = tapes[0]
+
+        assert new_tape.measurements == [
+            qml.expval(X(0)),
+            qml.var(X(1) + Z(2) @ X(1)),
+            qml.counts(X(0) @ (2.3 * Z(2))),
+        ]
+        assert new_tape.operations == Y(2).diagonalizing_gates()
+
+        assert fn == null_postprocessing
+
+    @pytest.mark.parametrize(
+        "supported_base_obs", (["PauliC", "PauliZ"], ["PauliX", "PauliZ"], [X(0), qml.Z(1)])
+    )
+    def test_bad_obs_input_raises_error(self, supported_base_obs):
+        """Test that if a value is passed to supported_base_obs that can't be interpreted, a clear error is raised"""
+
+        with pytest.raises(ValueError, match="Supported base observables must be a subset of"):
+            _ = diagonalize_measurements(
+                QuantumScript([], measurements=[]), supported_base_obs=supported_base_obs
+            )
+
+    @pytest.mark.parametrize("supported_base_obs", ([qml.Z], [qml.Z, qml.X], [qml.Z, qml.X, qml.Y]))
+    @pytest.mark.parametrize("shots", [None, 2000, (4000, 5000, 6000)])
+    def test_qnode_integration(self, supported_base_obs, shots):
+
+        dev = qml.device("default.qubit", shots=shots)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.RX(1.23, 1)
+            qml.RY(2.46, 0)
+            return qml.expval(X(0)), qml.var(X(1) + Y(2))
+
+        @partial(diagonalize_measurements, supported_base_obs=supported_base_obs)
+        @qml.qnode(dev)
+        def circuit_diagonalized():
+            qml.RX(1.23, 1)
+            qml.RY(2.46, 0)
+            return qml.expval(X(0)), qml.var(X(1) + Y(2))
+
+        expected_res = circuit()
+        res = circuit_diagonalized()
+
+        if len(dev.shots.shot_vector) > 1:
+            for r_diagonalized, r in zip(res, expected_res):
+                assert np.allclose(r_diagonalized, r, atol=0.1)
+        else:
+            assert np.allclose(expected_res, res, atol=0.1)