Fix transforms that error with non-commuting measurements (#5424)

**Context:**

`tape.expand` will error out if the tape has any non-commuting
measurements present. But multiple transforms that don't care about
measurements at all use it to decompose to a target gate set. Even when
`expand_measurements=False`, `tape.expand` requires measurements to be
commuting, or composite observables. And this is a core behavior of
`tape.expand` that people might rely on. It's not as much a bug as a
problematic feature. AKA load-bearing bug.

Fortunately, we now have `qml.devices.preprocess.decompose`, which is a
transform that is capable of decomposing to a target gateset. While yes,
this transform is in the `devices` module, so using it is introducing a
cross-module dependency in a direction that we do not want, it gets the
job done. I would call this a sign that we should potentially consider
moving it into the `transforms` module and making it's import path more
accessible.

**Description of the Change:**

Use `decompose` instead of `tape.expand` in a variety of transforms.

**Benefits:**

These transforms now allow non-commuting measurements through.

**Possible Drawbacks:**

We still have other transforms still relying on `tape.expand`. But we
can investigate those another day.

Introducing a cross-module dependency in a direction we don't want a
dependency. This could be fixed by moving `decompose` to the transforms
module, but that is a problem for another day.

**Related GitHub Issues:**

Fixes #5401 Fixes #5316 [sc-58086]  [sc-59118]

---------

Co-authored-by: Mudit Pandey <mudit.pandey@xanadu.ai>
Co-authored-by: Astral Cai <astral.cai@xanadu.ai>
Co-authored-by: Vincent Michaud-Rioux <vincentm@nanoacademic.com>
Co-authored-by: David Wierichs <david.wierichs@xanadu.ai>
Co-authored-by: lillian542 <38584660+lillian542@users.noreply.github.com>

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):

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

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 decomposer(op):
     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
 

pennylane/gradients/finite_difference.py

@@ -30,7 +30,6 @@
 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(

pennylane/gradients/parameter_shift.py

@@ -26,7 +26,6 @@
 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(

pennylane/ops/op_math/pow.py

@@ -262,6 +262,11 @@ def has_decomposition(self):
             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 @@ def decomposition(self):
             # 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):

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 qfunc(x, y, z):
         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):

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

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 qfunc(x, y, z):
         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:

pennylane/transforms/transpile.py

@@ -161,10 +161,15 @@ def circuit():
     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

tests/devices/default_qubit/test_default_qubit_preprocessing.py

@@ -862,12 +862,12 @@ def test_valid_tape_no_expand(self, G):
             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

tests/devices/test_default_clifford.py

@@ -632,7 +632,7 @@ def circuit():
 
     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 circ_2():
 
     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()
 

tests/devices/test_preprocess.py

@@ -115,7 +115,7 @@ def test_error_from_unsupported_operation(self):
         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 @@ def test_error_if_invalid_op(self):
         """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):

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 @@ def test_correct_second_derivative_center_order4(self):
 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."""