Repository: PennyLaneAI/pennylane

Issue:
[Bug]  `qml.transforms.merge_rotations` cannot handle non-commuting observables
Related to https://github.com/PennyLaneAI/pennylane/issues/5316
Came up in https://github.com/PennyLaneAI/pennylane/pull/5396

```python
@qml.transforms.merge_rotations
@qml.qnode(qml.device("default.qubit.legacy", wires=2, shots=50))
def circuit(order):
    qml.Permute(order, wires=(0, 1, 2))
    qml.RX(0.5, wires=0)
    qml.RX(-0.5, wires=0)
    return [qml.expval(qml.PauliX(0)), qml.expval(qml.PauliY(0))]

batch, fn = qml.workflow.construct_batch(circuit, level=None)((2, 1, 0))
```

```
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
File ~/Xanadu/pennylane/pennylane/tape/tape.py:88, in rotations_and_diagonal_measurements(tape)
     87 try:
---> 88     rotations, diag_obs = qml.pauli.diagonalize_qwc_pauli_words(tape._obs_sharing_wires)
     89 except (TypeError, ValueError) as e:

File ~/anaconda3/envs/pennylane311/lib/python3.11/contextlib.py:81, in ContextDecorator.__call__.<locals>.inner(*args, **kwds)
     80 with self._recreate_cm():
---> 81     return func(*args, **kwds)

File ~/Xanadu/pennylane/pennylane/pauli/utils.py:1142, in diagonalize_qwc_pauli_words(qwc_grouping)
   1141     if full_pauli_word[wire] != pauli_type:
-> 1142         raise ValueError("The list of Pauli words are not qubit-wise commuting")
   1143 else:

ValueError: The list of Pauli words are not qubit-wise commuting

The above exception was the direct cause of the following exception:

QuantumFunctionError                      Traceback (most recent call last)
Cell In[6], line 9
      6     qml.RX(-0.5, wires=0)
      7     return [qml.expval(qml.PauliX(0)), qml.expval(qml.PauliY(0))]
----> 9 batch, fn = qml.workflow.construct_batch(circuit, level=None)((2, 1, 0))

File ~/Xanadu/pennylane/pennylane/workflow/construct_batch.py:302, in construct_batch.<locals>.batch_constructor(*args, **kwargs)
    300 qnode._update_gradient_fn(tape=initial_tape)
    301 program = get_transform_program(qnode, level=level)
--> 302 return program((initial_tape,))

File ~/Xanadu/pennylane/pennylane/transforms/core/transform_program.py:509, in TransformProgram.__call__(self, tapes)
    507 if self._argnums is not None and self._argnums[i] is not None:
    508     tape.trainable_params = self._argnums[i][j]
--> 509 new_tapes, fn = transform(tape, *targs, **tkwargs)
    510 execution_tapes.extend(new_tapes)
    512 fns.append(fn)

File ~/Xanadu/pennylane/pennylane/transforms/optimization/merge_rotations.py:120, in merge_rotations(tape, atol, include_gates)
     32 r"""Quantum transform to combine rotation gates of the same type that act sequentially.
     33 
     34 If the combination of two rotation produces an angle that is close to 0,
   (...)
    117 
    118 """
    119 # Expand away adjoint ops
--> 120 expanded_tape = tape.expand(stop_at=lambda obj: not isinstance(obj, Adjoint))
    121 list_copy = expanded_tape.operations
    122 new_operations = []

File ~/Xanadu/pennylane/pennylane/tape/qscript.py:912, in QuantumScript.expand(self, depth, stop_at, expand_measurements)
    867 def expand(self, depth=1, stop_at=None, expand_measurements=False):
    868     """Expand all operations to a specific depth.
    869 
    870     Args:
   (...)
    910     RY(0.2, wires=['a'])]
    911     """
--> 912     new_script = qml.tape.tape.expand_tape(
    913         self, depth=depth, stop_at=stop_at, expand_measurements=expand_measurements
    914     )
    915     new_script._update()
    916     return new_script

File ~/Xanadu/pennylane/pennylane/tape/tape.py:181, in expand_tape(tape, depth, stop_at, expand_measurements)
    178 if tape.samples_computational_basis and len(tape.measurements) > 1:
    179     _validate_computational_basis_sampling(tape.measurements)
--> 181 diagonalizing_gates, diagonal_measurements = rotations_and_diagonal_measurements(tape)
    182 for queue, new_queue in [
    183     (tape.operations + diagonalizing_gates, new_ops),
    184     (diagonal_measurements, new_measurements),
    185 ]:
    186     for obj in queue:

File ~/Xanadu/pennylane/pennylane/tape/tape.py:99, in rotations_and_diagonal_measurements(tape)
     90     if any(isinstance(m, (ProbabilityMP, SampleMP, CountsMP)) for m in tape.measurements):
     91         raise qml.QuantumFunctionError(
     92             "Only observables that are qubit-wise commuting "
     93             "Pauli words can be returned on the same wire.\n"
   (...)
     96             "for each non-commuting observable."
     97         ) from e
---> 99     raise qml.QuantumFunctionError(_err_msg_for_some_meas_not_qwc(tape.measurements)) from e
    101 measurements = copy.copy(tape.measurements)
    103 for o, i in zip(diag_obs, tape._obs_sharing_wires_id):

QuantumFunctionError: Only observables that are qubit-wise commuting Pauli words can be returned on the same wire, some of the following measurements do not commute:
[expval(X(0)), expval(Y(0))]
```

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5424?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 `89.79592%` with `5 lines` in your changes are missing coverage. Please review.
> :exclamation: No coverage uploaded for pull request base (`v0.36.0-rc0@ab5785f`). [Click here to learn what that means](https://docs.codecov.io/docs/error-reference?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#section-missing-base-commit).

| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5424?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/ops/op\_math/pow.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5424?src=pr&el=tree&filepath=pennylane%2Fops%2Fop_math%2Fpow.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9vcF9tYXRoL3Bvdy5weQ==) | 50.00% | [3 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5424?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |
| [pennylane/transforms/compile.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5424?src=pr&el=tree&filepath=pennylane%2Ftransforms%2Fcompile.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvY29tcGlsZS5weQ==) | 75.00% | [1 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5424?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |
| [pennylane/transforms/transpile.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5424?src=pr&el=tree&filepath=pennylane%2Ftransforms%2Ftranspile.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvdHJhbnNwaWxlLnB5) | 66.66% | [1 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5424?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               @@
##             v0.36.0-rc0    #5424   +/-   ##
==============================================
  Coverage               ?   99.67%           
==============================================
  Files                  ?      412           
  Lines                  ?    38313           
  Branches               ?        0           
==============================================
  Hits                   ?    38188           
  Misses                 ?      125           
  Partials               ?        0           
```



</details>

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


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index ad7395f90..ea113ec99 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -71,6 +71,10 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* `param_shift`, `finite_diff`, `compile`, `merge_rotations`, and `transpile` now all work
+  with circuits with non-commuting measurements.
+  [(#5424)](https://github.com/PennyLaneAI/pennylane/pull/5424)
+
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):
diff --git a/pennylane/devices/default_clifford.py b/pennylane/devices/default_clifford.py
index 3ea29fd64..7a1d31035 100644
--- a/pennylane/devices/default_clifford.py
+++ b/pennylane/devices/default_clifford.py
@@ -136,7 +136,7 @@ def _pl_op_to_stim(op):
         stim_tg = map(str, op.wires)
     except KeyError as e:
         raise qml.DeviceError(
-            f"Operator {op} not supported on default.clifford and does not provide a decomposition."
+            f"Operator {op} not supported with default.clifford and does not provide a decomposition."
         ) from e
 
     # Check if the operation is noisy
diff --git a/pennylane/devices/preprocess.py b/pennylane/devices/preprocess.py
index 2c821306a..98d6442f4 100644
--- a/pennylane/devices/preprocess.py
+++ b/pennylane/devices/preprocess.py
@@ -58,7 +58,7 @@ def _operator_decomposition_gen(
             current_depth += 1
         except qml.operation.DecompositionUndefinedError as e:
             raise DeviceError(
-                f"Operator {op} not supported on {name} and does not provide a decomposition."
+                f"Operator {op} not supported with {name} and does not provide a decomposition."
             ) from e
 
         for sub_op in decomp:
@@ -325,30 +325,32 @@ def decompose(
     if stopping_condition_shots is not None and tape.shots:
         stopping_condition = stopping_condition_shots
 
-    if not all(stopping_condition(op) for op in tape.operations):
-        try:
-            # don't decompose initial operations if its StatePrepBase
-            prep_op = (
-                [tape[0]] if isinstance(tape[0], StatePrepBase) and skip_initial_state_prep else []
+    if tape.operations and isinstance(tape[0], StatePrepBase) and skip_initial_state_prep:
+        prep_op = [tape[0]]
+    else:
+        prep_op = []
+
+    if all(stopping_condition(op) for op in tape.operations[len(prep_op) :]):
+        return (tape,), null_postprocessing
+    try:
+
+        new_ops = [
+            final_op
+            for op in tape.operations[len(prep_op) :]
+            for final_op in _operator_decomposition_gen(
+                op,
+                stopping_condition,
+                decomposer=decomposer,
+                max_expansion=max_expansion,
+                name=name,
             )
-
-            new_ops = [
-                final_op
-                for op in tape.operations[bool(prep_op) :]
-                for final_op in _operator_decomposition_gen(
-                    op,
-                    stopping_condition,
-                    decomposer=decomposer,
-                    max_expansion=max_expansion,
-                    name=name,
-                )
-            ]
-        except RecursionError as e:
-            raise DeviceError(
-                "Reached recursion limit trying to decompose operations. "
-                "Operator decomposition may have entered an infinite loop."
-            ) from e
-        tape = qml.tape.QuantumScript(prep_op + new_ops, tape.measurements, shots=tape.shots)
+        ]
+    except RecursionError as e:
+        raise DeviceError(
+            "Reached recursion limit trying to decompose operations. "
+            "Operator decomposition may have entered an infinite loop."
+        ) from e
+    tape = qml.tape.QuantumScript(prep_op + new_ops, tape.measurements, shots=tape.shots)
 
     return (tape,), null_postprocessing
 
diff --git a/pennylane/gradients/finite_difference.py b/pennylane/gradients/finite_difference.py
index 1d7e7e8e7..e478bc653 100644
--- a/pennylane/gradients/finite_difference.py
+++ b/pennylane/gradients/finite_difference.py
@@ -30,7 +30,6 @@ import pennylane as qml
 from pennylane import transform
 from pennylane.gradients.gradient_transform import _contract_qjac_with_cjac
 from pennylane.measurements import ProbabilityMP
-from pennylane.transforms.tape_expand import expand_invalid_trainable
 
 from .general_shift_rules import generate_shifted_tapes
 from .gradient_transform import (
@@ -177,6 +176,17 @@ def _processing_fn(results, shots, single_shot_batch_fn):
     return tuple(grads_tuple)
 
 
+def _finite_diff_stopping_condition(op) -> bool:
+    if not op.has_decomposition:
+        # let things without decompositions through without error
+        # error will happen when calculating shifted tapes for finite diff
+
+        return True
+    if isinstance(op, qml.operation.Operator) and any(qml.math.requires_grad(p) for p in op.data):
+        return op.grad_method is not None
+    return True
+
+
 def _expand_transform_finite_diff(
     tape: qml.tape.QuantumTape,
     argnum=None,
@@ -188,15 +198,17 @@ def _expand_transform_finite_diff(
     validate_params=True,
 ) -> (Sequence[qml.tape.QuantumTape], Callable):
     """Expand function to be applied before finite difference."""
-    expanded_tape = expand_invalid_trainable(tape)
-
-    def null_postprocessing(results):
-        """A postprocesing function returned by a transform that only converts the batch of results
-        into a result for a single ``QuantumTape``.
-        """
-        return results[0]
-
-    return [expanded_tape], null_postprocessing
+    [new_tape], postprocessing = qml.devices.preprocess.decompose(
+        tape,
+        stopping_condition=_finite_diff_stopping_condition,
+        skip_initial_state_prep=False,
+        name="finite_diff",
+    )
+    if new_tape is tape:
+        return [tape], postprocessing
+    params = new_tape.get_parameters(trainable_only=False)
+    new_tape.trainable_params = qml.math.get_trainable_indices(params)
+    return [new_tape], postprocessing
 
 
 @partial(
diff --git a/pennylane/gradients/parameter_shift.py b/pennylane/gradients/parameter_shift.py
index 9f48171cd..ef2709e7a 100644
--- a/pennylane/gradients/parameter_shift.py
+++ b/pennylane/gradients/parameter_shift.py
@@ -26,7 +26,6 @@ import pennylane as qml
 from pennylane import transform
 from pennylane.gradients.gradient_transform import _contract_qjac_with_cjac
 from pennylane.measurements import VarianceMP
-from pennylane.transforms.tape_expand import expand_invalid_trainable
 
 from .finite_difference import finite_diff
 from .general_shift_rules import (
@@ -756,6 +755,16 @@ def var_param_shift(tape, argnum, shifts=None, gradient_recipes=None, f0=None, b
     return gradient_tapes, processing_fn
 
 
+def _param_shift_stopping_condition(op) -> bool:
+    if not op.has_decomposition:
+        # let things without decompositions through without error
+        # error will happen when calculating parameter shift tapes
+        return True
+    if isinstance(op, qml.operation.Operator) and any(qml.math.requires_grad(p) for p in op.data):
+        return op.grad_method is not None
+    return True
+
+
 def _expand_transform_param_shift(
     tape: qml.tape.QuantumTape,
     argnum=None,
@@ -766,15 +775,17 @@ def _expand_transform_param_shift(
     broadcast=False,
 ) -> (Sequence[qml.tape.QuantumTape], Callable):
     """Expand function to be applied before parameter shift."""
-    expanded_tape = expand_invalid_trainable(tape)
-
-    def null_postprocessing(results):
-        """A postprocesing function returned by a transform that only converts the batch of results
-        into a result for a single ``QuantumTape``.
-        """
-        return results[0]
-
-    return [expanded_tape], null_postprocessing
+    [new_tape], postprocessing = qml.devices.preprocess.decompose(
+        tape,
+        stopping_condition=_param_shift_stopping_condition,
+        skip_initial_state_prep=False,
+        name="param_shift",
+    )
+    if new_tape is tape:
+        return [tape], postprocessing
+    params = new_tape.get_parameters(trainable_only=False)
+    new_tape.trainable_params = qml.math.get_trainable_indices(params)
+    return [new_tape], postprocessing
 
 
 @partial(
diff --git a/pennylane/ops/op_math/pow.py b/pennylane/ops/op_math/pow.py
index 30b214c2a..98a724b06 100644
--- a/pennylane/ops/op_math/pow.py
+++ b/pennylane/ops/op_math/pow.py
@@ -262,6 +262,11 @@ class Pow(ScalarSymbolicOp):
             self.base.pow(self.z)
         except PowUndefinedError:
             return False
+        except Exception as e:  # pylint: disable=broad-except
+            # some pow methods cant handle a batched z
+            if qml.math.ndim(self.z) != 0:
+                return False
+            raise e
         return True
 
     def decomposition(self):
@@ -275,6 +280,8 @@ class Pow(ScalarSymbolicOp):
             # TODO: consider: what if z is an int and less than 0?
             # do we want Pow(base, -1) to be a "more fundamental" op
             raise DecompositionUndefinedError from e
+        except Exception as e:  # pylint: disable=broad-except
+            raise DecompositionUndefinedError from e
 
     @property
     def has_diagonalizing_gates(self):
diff --git a/pennylane/transforms/compile.py b/pennylane/transforms/compile.py
index 696d330cd..a2421a984 100644
--- a/pennylane/transforms/compile.py
+++ b/pennylane/transforms/compile.py
@@ -16,6 +16,7 @@
 from functools import partial
 from typing import Callable, Sequence
 
+import pennylane as qml
 from pennylane.ops import __all__ as all_ops
 from pennylane.queuing import QueuingManager
 from pennylane.tape import QuantumTape
@@ -182,9 +183,15 @@ def compile(
         basis_set = basis_set or all_ops
 
         def stop_at(obj):
+            if not isinstance(obj, qml.operation.Operator):
+                return True
+            if not obj.has_decomposition:
+                return True
             return obj.name in basis_set and (not getattr(obj, "only_visual", False))
 
-        expanded_tape = tape.expand(depth=expand_depth, stop_at=stop_at)
+        [expanded_tape], _ = qml.devices.preprocess.decompose(
+            tape, stopping_condition=stop_at, max_expansion=expand_depth, name="compile"
+        )
 
         # Apply the full set of compilation transforms num_passes times
         for _ in range(num_passes):
diff --git a/pennylane/transforms/optimization/commute_controlled.py b/pennylane/transforms/optimization/commute_controlled.py
index 2b22535c3..9314b051b 100644
--- a/pennylane/transforms/optimization/commute_controlled.py
+++ b/pennylane/transforms/optimization/commute_controlled.py
@@ -42,7 +42,7 @@ def _commute_controlled_right(op_list):
         # We are looking only at the gates that can be pushed through
         # controls/targets; these are single-qubit gates with the basis
         # property specified.
-        if current_gate.basis is None or len(current_gate.wires) != 1:
+        if getattr(current_gate, "basis", None) is None or len(current_gate.wires) != 1:
             current_location -= 1
             continue
 
@@ -56,7 +56,7 @@ def _commute_controlled_right(op_list):
             next_gate = op_list[new_location + next_gate_idx + 1]
 
             # Only go ahead if information is available
-            if next_gate.basis is None:
+            if getattr(next_gate, "basis", None) is None:
                 break
 
             # If the next gate does not have control_wires defined, it is not
diff --git a/pennylane/transforms/optimization/merge_rotations.py b/pennylane/transforms/optimization/merge_rotations.py
index 6b19f8cf8..e3a7362ae 100644
--- a/pennylane/transforms/optimization/merge_rotations.py
+++ b/pennylane/transforms/optimization/merge_rotations.py
@@ -15,6 +15,7 @@
 # pylint: disable=too-many-branches
 from typing import Callable, Sequence
 
+import pennylane as qml
 from pennylane.math import allclose, cast_like, get_interface, is_abstract, stack, zeros
 from pennylane.ops.op_math import Adjoint
 from pennylane.ops.qubit.attributes import composable_rotations
@@ -116,8 +117,14 @@ def merge_rotations(
         2: ─╰X─────────H────────╰●───────────────────┤
 
     """
+
     # Expand away adjoint ops
-    expanded_tape = tape.expand(stop_at=lambda obj: not isinstance(obj, Adjoint))
+    def stop_at(obj):
+        return not isinstance(obj, Adjoint)
+
+    [expanded_tape], _ = qml.devices.preprocess.decompose(
+        tape, stopping_condition=stop_at, name="merge_rotations"
+    )
     list_copy = expanded_tape.operations
     new_operations = []
     while len(list_copy) > 0:
diff --git a/pennylane/transforms/transpile.py b/pennylane/transforms/transpile.py
index 12e0adba1..f2f6e4e08 100644
--- a/pennylane/transforms/transpile.py
+++ b/pennylane/transforms/transpile.py
@@ -161,10 +161,15 @@ def transpile(
     with QueuingManager.stop_recording():
         # this unrolls everything in the current tape (in particular templates)
         def stop_at(obj):
+            if not isinstance(obj, qml.operation.Operator):
+                return True
+            if not obj.has_decomposition:
+                return True
             return (obj.name in all_ops) and (not getattr(obj, "only_visual", False))
 
-        expanded_tape = tape.expand(stop_at=stop_at)
-
+        [expanded_tape], _ = qml.devices.preprocess.decompose(
+            tape, stopping_condition=stop_at, name="transpile"
+        )
         # make copy of ops
         list_op_copy = expanded_tape.operations.copy()
         wire_order = device_wires or tape.wires


Test patch:
diff --git a/tests/devices/default_qubit/test_default_qubit_preprocessing.py b/tests/devices/default_qubit/test_default_qubit_preprocessing.py
index ac5e80d2b..1b7a72d49 100644
--- a/tests/devices/default_qubit/test_default_qubit_preprocessing.py
+++ b/tests/devices/default_qubit/test_default_qubit_preprocessing.py
@@ -862,12 +862,12 @@ class TestAdjointDiffTapeValidation:
             ExecutionConfig(gradient_method="adjoint")
         )[0]
 
-        qs.trainable_params = {1}
+        qs.trainable_params = [1]
         qs_valid, _ = program((qs,))
         qs_valid = qs_valid[0]
         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_valid.trainable_params == [0, 1]
+        assert qs_valid.trainable_params == [1]  # same as input tape since no decomposition
 
     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
diff --git a/tests/devices/test_default_clifford.py b/tests/devices/test_default_clifford.py
index 9df6dc25d..9178b6a14 100644
--- a/tests/devices/test_default_clifford.py
+++ b/tests/devices/test_default_clifford.py
@@ -632,7 +632,7 @@ def test_clifford_error(check):
 
     with pytest.raises(
         qml.DeviceError,
-        match=r"Operator RX\(1.0, wires=\[0\]\) not supported on default.clifford and does not provide a decomposition",
+        match=r"Operator RX\(1.0, wires=\[0\]\) not supported with default.clifford and does not provide a decomposition",
     ):
         circuit()
 
@@ -660,7 +660,7 @@ def test_meas_error_noisy():
 
     with pytest.raises(
         qml.DeviceError,
-        match=r"Operator AmplitudeDamping\(0.2, wires=\[0\]\) not supported on default.clifford",
+        match=r"Operator AmplitudeDamping\(0.2, wires=\[0\]\) not supported with default.clifford",
     ):
         circ_2()
 
diff --git a/tests/devices/test_preprocess.py b/tests/devices/test_preprocess.py
index 0931a8d99..fb5897041 100644
--- a/tests/devices/test_preprocess.py
+++ b/tests/devices/test_preprocess.py
@@ -115,7 +115,7 @@ class TestPrivateHelpers:
         op = NoMatNoDecompOp("a")
         with pytest.raises(
             DeviceError,
-            match=r"not supported on abc and does",
+            match=r"not supported with abc and does",
         ):
             tuple(
                 _operator_decomposition_gen(
@@ -192,7 +192,7 @@ class TestDecomposeValidation:
         """Test that expand_fn throws an error when an operation is does not define a matrix or decomposition."""
 
         tape = QuantumScript(ops=[NoMatNoDecompOp(0)], measurements=[qml.expval(qml.Hadamard(0))])
-        with pytest.raises(DeviceError, match="not supported on abc"):
+        with pytest.raises(DeviceError, match="not supported with abc"):
             decompose(tape, lambda op: op.has_matrix, name="abc")
 
     def test_decompose(self):
diff --git a/tests/gradients/finite_diff/test_finite_difference.py b/tests/gradients/finite_diff/test_finite_difference.py
index d135382b3..4031c7ea4 100644
--- a/tests/gradients/finite_diff/test_finite_difference.py
+++ b/tests/gradients/finite_diff/test_finite_difference.py
@@ -14,6 +14,7 @@
 """
 Tests for the gradients.finite_difference module.
 """
+# pylint: disable=use-implicit-booleaness-not-comparison
 import numpy
 import pytest
 
@@ -99,6 +100,20 @@ class TestCoeffs:
 class TestFiniteDiff:
     """Tests for the finite difference gradient transform"""
 
+    def test_finite_diff_non_commuting_observables(self):
+        """Test that finite differences work even if the measurements do not commute with each other."""
+
+        ops = (qml.RX(0.5, wires=0),)
+        ms = (qml.expval(qml.X(0)), qml.expval(qml.Z(0)))
+        tape = qml.tape.QuantumScript(ops, ms, trainable_params=[0])
+
+        batch, _ = qml.gradients.finite_diff(tape)
+        assert len(batch) == 2
+        tape0 = qml.tape.QuantumScript((qml.RX(0.5, 0),), ms, trainable_params=[0])
+        tape1 = qml.tape.QuantumScript((qml.RX(0.5 + 1e-7, 0),), ms, trainable_params=[0])
+        assert qml.equal(batch[0], tape0)
+        assert qml.equal(batch[1], tape1)
+
     def test_trainable_batched_tape_raises(self):
         """Test that an error is raised for a broadcasted/batched tape if the broadcasted
         parameter is differentiated."""
diff --git a/tests/gradients/parameter_shift/test_parameter_shift.py b/tests/gradients/parameter_shift/test_parameter_shift.py
index ef982b0ba..4a33f631b 100644
--- a/tests/gradients/parameter_shift/test_parameter_shift.py
+++ b/tests/gradients/parameter_shift/test_parameter_shift.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Tests for the gradients.parameter_shift module using the new return types."""
+# pylint: disable=use-implicit-booleaness-not-comparison
 import pytest
 
 import pennylane as qml
@@ -341,6 +342,20 @@ class TestParamShift:
         tapes, _ = qml.gradients.param_shift(tape)
         assert not tapes
 
+    def test_parameter_shift_non_commuting_observables(self):
+        """Test that parameter shift works even if the measurements do not commute with each other."""
+
+        ops = (qml.RX(0.5, wires=0),)
+        ms = (qml.expval(qml.X(0)), qml.expval(qml.Z(0)))
+        tape = qml.tape.QuantumScript(ops, ms, trainable_params=[0])
+
+        batch, _ = qml.gradients.param_shift(tape)
+        assert len(batch) == 2
+        tape0 = qml.tape.QuantumScript((qml.RX(0.5 + np.pi / 2, 0),), ms, trainable_params=[0])
+        tape1 = qml.tape.QuantumScript((qml.RX(0.5 - np.pi / 2, 0),), ms, trainable_params=[0])
+        assert qml.equal(batch[0], tape0)
+        assert qml.equal(batch[1], tape1)
+
     def test_state_non_differentiable_error(self):
         """Test error raised if attempting to differentiate with
         respect to a state"""
diff --git a/tests/ops/op_math/test_pow_op.py b/tests/ops/op_math/test_pow_op.py
index 68850f1f7..c1bbc07dc 100644
--- a/tests/ops/op_math/test_pow_op.py
+++ b/tests/ops/op_math/test_pow_op.py
@@ -343,6 +343,36 @@ class TestProperties:
 
         assert op.has_decomposition is False
 
+    def test_no_decomposition_batching_error(self, power_method):
+        """Test that if an error occurs with a batched exponent, has_decomposition is False."""
+
+        class MyOp(qml.operation.Operator):
+
+            def pow(self, z):
+                return super().pow(z % 2)
+
+        pow_op = power_method(base=MyOp(wires=0), z=np.array([1.0, 2.0]))
+        assert not pow_op.has_decomposition
+
+        with pytest.raises(DecompositionUndefinedError):
+            pow_op.decomposition()
+
+    def test_error_raised_if_no_batching(self, power_method):
+        """Test that if Operator.pow raises an error and no batching is present, the erorr is raised."""
+
+        class MyOp(qml.operation.Operator):
+
+            def pow(self, z):
+                raise ValueError
+
+        pow_op = power_method(base=MyOp(0), z=2.5)
+
+        with pytest.raises(ValueError):
+            _ = pow_op.has_decomposition
+
+        with pytest.raises(DecompositionUndefinedError):
+            _ = pow_op.decomposition()
+
     @pytest.mark.parametrize("value", (True, False))
     def test_has_diagonalizing_gates(self, value, power_method):
         """Test that Pow defers has_diagonalizing_gates to base operator."""
diff --git a/tests/transforms/test_compile.py b/tests/transforms/test_compile.py
index f7f26387d..599c7336d 100644
--- a/tests/transforms/test_compile.py
+++ b/tests/transforms/test_compile.py
@@ -99,6 +99,28 @@ class TestCompile:
 
         compare_operation_lists(transformed_qnode.qtape.operations, names_expected, wires_expected)
 
+    def test_compile_non_commuting_observables(self):
+        """Test that compile works with non-commuting observables."""
+
+        ops = (qml.RX(0.1, 0), qml.RX(0.2, 0), qml.Barrier(only_visual=True), qml.X(0), qml.X(0))
+        ms = (qml.expval(qml.X(0)), qml.expval(qml.Y(0)))
+        tape = qml.tape.QuantumScript(ops, ms, shots=50)
+
+        [out], _ = qml.compile(tape)
+        expected = qml.tape.QuantumScript((qml.RX(0.3, 0),), ms, shots=50)
+        assert qml.equal(out, expected)
+
+    def test_compile_mcm(self):
+        """Test that compile works with mid circuit measurements and conditionals."""
+
+        m0 = qml.measure(0)
+        ops = [qml.X(0), qml.X(0), *m0.measurements, qml.ops.Conditional(m0, qml.S(0))]
+        tape = qml.tape.QuantumScript(ops, [qml.probs()], shots=50)
+
+        [new_tape], _ = qml.compile(tape)
+        assert new_tape.shots == tape.shots
+        assert new_tape.circuit == tape[2:]
+
 
 class TestCompileIntegration:
     """Integration tests to verify outputs of compilation pipelines."""
@@ -479,8 +501,8 @@ class TestCompileInterfaces:
 
         # Check that the gradient is the same
         assert qml.math.allclose(
-            jax.grad(original_qnode, argnums=(1))(x, params),
-            jax.grad(transformed_qnode, argnums=(1))(x, params),
+            jax.grad(original_qnode, argnums=1)(x, params),
+            jax.grad(transformed_qnode, argnums=1)(x, params),
             atol=1e-7,
         )
 
diff --git a/tests/transforms/test_optimization/test_merge_rotations.py b/tests/transforms/test_optimization/test_merge_rotations.py
index fa2333399..82453cb62 100644
--- a/tests/transforms/test_optimization/test_merge_rotations.py
+++ b/tests/transforms/test_optimization/test_merge_rotations.py
@@ -261,6 +261,17 @@ class TestMergeRotations:
         assert ops[0].parameters[0] == 0.2
         assert ops[1].parameters[0] == 0.3
 
+    def test_merge_rotations_non_commuting_observables(self):
+        """Test that merge_rotations can be used with non-commuting observables."""
+
+        ops = (qml.RX(0.1, 0), qml.RX(0.3, 0))
+        ms = (qml.expval(qml.X(0)), qml.expval(qml.Y(0)))
+
+        tape = qml.tape.QuantumScript(ops, ms, shots=50)
+        [out], _ = qml.transforms.merge_rotations(tape)
+        expected = qml.tape.QuantumScript((qml.RX(0.4, 0),), ms, shots=50)
+        assert qml.equal(out, expected)
+
 
 # Example QNode and device for interface testing
 dev = qml.device("default.qubit", wires=3)
diff --git a/tests/transforms/test_transpile.py b/tests/transforms/test_transpile.py
index 0d72300b2..c97907c91 100644
--- a/tests/transforms/test_transpile.py
+++ b/tests/transforms/test_transpile.py
@@ -45,6 +45,7 @@ def build_qfunc_pauli_z(wires):
     return qfunc
 
 
+# pylint: disable=too-many-public-methods
 class TestTranspile:
     """Unit tests for transpile function"""
 
@@ -114,6 +115,17 @@ class TestTranspile:
         with pytest.raises(NotImplementedError, match=err_msg):
             transpiled_qnode()
 
+    def test_transpile_non_commuting_observables(self):
+        """Test that transpile will work with non-commuting observables."""
+
+        ops = (qml.CRX(0.1, wires=(0, 2)),)
+        ms = (qml.expval(qml.X(0)), qml.expval(qml.Y(0)))
+        tape = qml.tape.QuantumScript(ops, ms, shots=50)
+        [out], _ = qml.transforms.transpile(tape, coupling_map=((0, 1), (1, 2)))
+
+        expected = qml.tape.QuantumScript((qml.SWAP((1, 2)), qml.CRX(0.1, (0, 1))), ms, shots=50)
+        assert qml.equal(out, expected)
+
     def test_transpile_qfunc_transpiled_mmt_obs(self):
         """test that transpile does not alter output for expectation value of an observable"""
         dev = qml.device("default.qubit", wires=[0, 1, 2])
@@ -286,6 +298,18 @@ class TestTranspile:
             original_expectation, transpiled_expectation, atol=np.finfo(float).eps
         )
 
+    def test_transpile_mcm(self):
+        """Test that transpile can be used with mid circuit measurements."""
+
+        m0 = qml.measure(0)
+        ops = [qml.CNOT((0, 2)), *m0.measurements, qml.ops.Conditional(m0, qml.S(0))]
+        tape = qml.tape.QuantumScript(ops, [qml.probs()], shots=50)
+
+        [new_tape], _ = transpile(tape, [(0, 1), (1, 2)])
+        expected_ops = [qml.SWAP((1, 2)), qml.CNOT((0, 1))] + ops[1:]
+        assert new_tape.operations == expected_ops
+        assert new_tape.shots == tape.shots
+
     def test_transpile_ops_anywires_1_qubit_qnode(self):
         """test that transpile does not alter output for expectation value of an observable if the qfunc contains
         1-qubit operations with AnyWires defined for the operation"""