Repository: PennyLaneAI/pennylane
Pull request #3857 (2023-03-02 19:51:36)

allow circuits with non commuting observables in `qml.devices.qubit.preprocess`
On testing for #3767 , I found that `preprocess` would raise an error if we needed to decompose an operation but the circuit had non-commuting measurements.  Given that `qml.devices.qubit.simulate` can handle non-commuting measurements fine, we shouldn't be raising that error.

Given that the problems stems from `qml.tape.tape.expand_tape`, I stop using `expand_tape`.  Instead, we just iterate over the decompositions from the operations using a generator.

While changing the file I also:

* Combined `check_validity` into `expand_fn`
* Removed `_supported_observable`, as it didn't need to be its own function
* Renamed `tapes` to `circuits`
* Set `DefaultExecutionConfig` as the default value for `execution_config` instead of `None`
* Raise an error for non-state-based measurements (until we start supporting finite shots)
* Raise an error if we have more than one state prep operation
* Expand QFT and GroverOperator above certain hardcoded thresholds on the number of wires.
* Add type hints
* Rename `_stopping_condition` to `_accepted_operator`

codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3857?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3857](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3857?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (efa4641) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/6daabc6cd265f5112e426bedcb57f4197b7fbb7c?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (6daabc6) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master    #3857   +/-   ##
=======================================
  Coverage   99.71%   99.71%           
=======================================
  Files         338      338           
  Lines       29485    29490    +5     
=======================================
+ Hits        29400    29405    +5     
  Misses         85       85           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3857?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/devices/qubit/preprocess.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3857?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvcXViaXQvcHJlcHJvY2Vzcy5weQ==) | `100.00% <100.00%> (ø)` | |

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)

albi3ro: [sc-34554]

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index d786c0268..bcab4e0e5 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -25,6 +25,9 @@
 * `AdaptiveOptimizer` is updated to use non-default user-defined qnode arguments.
   [(#3765)](https://github.com/PennyLaneAI/pennylane/pull/3765)
 
+* `pennylane.devices.qubit.preprocess` now allows circuits with non-commuting observables.
+  [(#3857)](https://github.com/PennyLaneAI/pennylane/pull/3857)
+
 * When using Jax-jit with gradient transforms the trainable parameters are correctly set (instead of every parameter 
   to be set as trainable), and therefore the derivatives are computed more efficiently.
   [(#3697)](https://github.com/PennyLaneAI/pennylane/pull/3697)
diff --git a/pennylane/devices/qubit/__init__.py b/pennylane/devices/qubit/__init__.py
index 91368b5d5..7e93f236d 100644
--- a/pennylane/devices/qubit/__init__.py
+++ b/pennylane/devices/qubit/__init__.py
@@ -23,13 +23,16 @@ at your own discretion.
 
     create_initial_state
     apply_operation
+    measure
     sample_state
     preprocess
     simulate
 """
 
-from .apply_operation import apply_operation
 from .initialize_state import create_initial_state
+from .apply_operation import apply_operation
+from .measure import measure
 from .sampling import sample_state
+
 from .preprocess import preprocess
 from .simulate import simulate
diff --git a/pennylane/devices/qubit/preprocess.py b/pennylane/devices/qubit/preprocess.py
index fb8d40a1d..735baa7d4 100644
--- a/pennylane/devices/qubit/preprocess.py
+++ b/pennylane/devices/qubit/preprocess.py
@@ -14,13 +14,17 @@
 
 """This module contains functions for preprocessing `QuantumTape` objects to ensure
 that they are supported for execution by a device."""
+# pylint: disable=protected-access
+from typing import Generator, Callable, Tuple
 
 import pennylane as qml
 
-from pennylane.operation import Observable, Tensor
-from pennylane.measurements import MidMeasureMP
+from pennylane.operation import Tensor
+from pennylane.measurements import MidMeasureMP, StateMeasurement
 from pennylane import DeviceError
 
+from ..experimental import ExecutionConfig, DefaultExecutionConfig
+
 # Update observable list. Current list is same as supported observables for
 # default.qubit.
 _observables = {
@@ -40,30 +44,40 @@ _observables = {
     "Evolution",
 }
 
+### UTILITY FUNCTIONS FOR EXPANDING UNSUPPORTED OPERATIONS ###
 
-def _stopping_condition(op):
-    """Specify whether or not an Operator object is supported by the device"""
-    return getattr(op, "has_matrix", False)
 
+def _accepted_operator(op: qml.operation.Operator) -> bool:
+    """Specify whether or not an Operator object is supported by the device."""
+    if op.name == "QFT" and len(op.wires) >= 6:
+        return False
+    if op.name == "GroverOperator" and len(op.wires) >= 13:
+        return False
+    return op.has_matrix
 
-def _supports_observable(observable):
-    """Checks if an observable is supported by this device.
 
-    Args:
-        observable (type or str): observable to be checked
+def _operator_decomposition_gen(
+    op: qml.operation.Operator,
+) -> Generator[qml.operation.Operator, None, None]:
+    """A generator that yields the next operation that is accepted by DefaultQubit2."""
+    if _accepted_operator(op):
+        yield op
+    else:
+        try:
+            decomp = op.decomposition()
+        except qml.operation.DecompositionUndefinedError as e:
+            raise DeviceError(
+                f"Operator {op} not supported on DefaultQubit2. Must provide either a matrix or a decomposition."
+            ) from e
 
-    Returns:
-        bool: ``True`` iff supplied observable is supported
-    """
-    if isinstance(observable, type) and issubclass(observable, Observable):
-        observable = observable.__name__
-    if isinstance(observable, str):
-        return observable in _observables
+        for sub_op in decomp:
+            yield from _operator_decomposition_gen(sub_op)
 
-    return False
 
+#######################
 
-def expand_fn(circuit, max_expansion=10):
+
+def expand_fn(circuit: qml.tape.QuantumScript) -> qml.tape.QuantumScript:
     """Method for expanding or decomposing an input circuit.
 
     This method expands the tape if:
@@ -73,11 +87,6 @@ def expand_fn(circuit, max_expansion=10):
 
     Args:
         circuit (.QuantumTape): the circuit to expand.
-        max_expansion (int): The number of times the circuit should be
-            expanded. Expansion occurs when an operation or measurement is not
-            supported, and results in a gate decomposition. If any operations
-            in the decomposition remain unsupported by the device, another
-            expansion occurs.
 
     Returns:
         .QuantumTape: The expanded/decomposed circuit, such that the device
@@ -87,13 +96,38 @@ def expand_fn(circuit, max_expansion=10):
     if any(isinstance(o, MidMeasureMP) for o in circuit.operations):
         circuit = qml.defer_measurements(circuit)
 
-    if not all(_stopping_condition(op) for op in circuit.operations):
-        circuit = circuit.expand(depth=max_expansion, stop_at=_stopping_condition)
+    if len(circuit._prep) > 1:
+        raise DeviceError("DefaultQubit2 accepts at most one state prep operation.")
+
+    if not all(_accepted_operator(op) for op in circuit._ops):
+        try:
+            new_ops = [
+                final_op for op in circuit._ops for final_op in _operator_decomposition_gen(op)
+            ]
+        except RecursionError as e:
+            raise DeviceError(
+                "Reached recursion limit trying to decompose operations. "
+                "Operator decomposition may have entered an infinite loop."
+            ) from e
+        circuit = qml.tape.QuantumScript(new_ops, circuit.measurements, circuit._prep)
+
+    for observable in circuit.observables:
+        if isinstance(observable, Tensor) and any(
+            o.name not in _observables for o in observable.obs
+        ):
+            raise DeviceError(f"Observable {observable} not supported on DefaultQubit2")
+        if observable.name not in _observables:
+            raise DeviceError(f"Observable {observable} not supported on DefaultQubit2")
+
+    # change this once shots are supported
+    for m in circuit.measurements:
+        if not isinstance(m, StateMeasurement):
+            raise DeviceError(f"Measurement process {m} is only useable with finite shots.")
 
     return circuit
 
 
-def batch_transform(circuit):
+def batch_transform(circuit: qml.tape.QuantumScript) -> (Tuple[qml.tape.QuantumScript], Callable):
     """Apply a differentiable batch transform for preprocessing a circuit
     prior to execution.
 
@@ -133,63 +167,31 @@ def batch_transform(circuit):
     return tapes, batch_fn
 
 
-def check_validity(tape):
-    """Checks whether the operations and observables in queue are all supported by the device.
-
-    Args:
-        tape (.QuantumTape): tape from which to validate operations and observables
-
-    Raises:
-        DeviceError: if there are operations in the queue or observables that the device does
-            not support
-    """
-    for o in tape.operations:
-        operation_name = o.name
-
-        if not _stopping_condition(o):
-            raise DeviceError(f"Gate {operation_name} not supported on Python Device")
-
-    for o in tape.observables:
-        if isinstance(o, Tensor):
-            for i in o.obs:
-                if not _supports_observable(i.name):
-                    raise DeviceError(f"Observable {i.name} not supported on Python Device")
-        else:
-            observable_name = o.name
-
-            if not _supports_observable(observable_name):
-                raise DeviceError(f"Observable {observable_name} not supported on Python Device")
-
-
-def preprocess(tapes, execution_config=None, max_expansion=10):
+def preprocess(
+    circuits: Tuple[qml.tape.QuantumScript],
+    execution_config: ExecutionConfig = DefaultExecutionConfig,
+) -> (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
     them to ensure all operators and measurements are supported by the execution device.
 
     Args:
-        tapes (Sequence[QuantumTape]): Batch of tapes to be processed.
+        circuits (Sequence[QuantumTape]): Batch of tapes to be processed.
         execution_config (.ExecutionConfig): execution configuration with configurable
             options for the execution.
-        max_expansion (int): The number of times the circuit should be
-            expanded. Expansion occurs when an operation or measurement is not
-            supported, and results in a gate decomposition. If any operations
-            in the decomposition remain unsupported by the device, another
-            expansion occurs.
 
     Returns:
         Tuple[Sequence[.QuantumTape], callable]: Returns a tuple containing
         the sequence of circuits to be executed, and a post-processing function
         to be applied to the list of evaluated circuit results.
     """
-    if execution_config and execution_config.shots is not None:
+    if execution_config.shots is not None:
         # Finite shot support will be added later
         raise DeviceError("The Python Device does not support finite shots.")
 
-    for i, tape in enumerate(tapes):
-        tapes[i] = expand_fn(tape, max_expansion=max_expansion)
-        check_validity(tapes[i])
+    circuits = tuple(expand_fn(c) for c in circuits)
 
-    tapes, batch_fn = qml.transforms.map_batch_transform(batch_transform, tapes)
+    circuits, batch_fn = qml.transforms.map_batch_transform(batch_transform, circuits)
 
-    return tapes, batch_fn
+    return circuits, batch_fn


Test patch:
diff --git a/tests/devices/qubit/test_preprocess.py b/tests/devices/qubit/test_preprocess.py
index 038771c7d..7af0abc54 100644
--- a/tests/devices/qubit/test_preprocess.py
+++ b/tests/devices/qubit/test_preprocess.py
@@ -19,10 +19,9 @@ import numpy as np
 import pennylane as qml
 from pennylane.operation import Operation
 from pennylane.devices.qubit.preprocess import (
-    _stopping_condition,
-    _supports_observable,
+    _accepted_operator,
+    _operator_decomposition_gen,
     expand_fn,
-    check_validity,
     batch_transform,
     preprocess,
 )
@@ -59,9 +58,7 @@ class NoMatNoDecompOp(Operation):
         return False
 
 
-class TestHelpers:
-    """Unit tests for helper functions in qml.devices.qubit.preprocess"""
-
+class TestPrivateHelpers:
     @pytest.mark.parametrize(
         "op, expected",
         [
@@ -69,54 +66,121 @@ class TestHelpers:
             (qml.CRX(0.1, wires=[0, 1]), True),
             (qml.Snapshot(), False),
             (qml.Barrier(), False),
+            (qml.QFT(wires=range(5)), True),
+            (qml.QFT(wires=range(10)), False),
+            (qml.GroverOperator(wires=range(10)), True),
+            (qml.GroverOperator(wires=range(14)), False),
         ],
     )
-    def test_stopping_condition(self, op, expected):
-        """Test that _stopping_condition works correctly"""
-        res = _stopping_condition(op)
+    def test_accepted_operator(self, op, expected):
+        """Test that _accepted_operator works correctly"""
+        res = _accepted_operator(op)
         assert res == expected
 
-    @pytest.mark.parametrize(
-        "obs, expected",
-        [("Hamiltonian", True), (qml.Identity, True), ("QubitUnitary", False), (qml.RX, False)],
-    )
-    def test_supports_observable(self, obs, expected):
-        """Test that _supports_observable works correctly"""
-        res = _supports_observable(obs)
-        assert res == expected
+    @pytest.mark.parametrize("op", (qml.PauliX(0), qml.RX(1.2, wires=0), qml.QFT(wires=range(3))))
+    def test_operator_decomposition_gen_accepted_operator(self, op):
+        casted_to_list = list(_operator_decomposition_gen(op))
+        assert len(casted_to_list) == 1
+        assert casted_to_list[0] is op
+
+    def test_operator_decomposition_gen_decomposed_operators_single_nesting(self):
+        """Assert _operator_decomposition_gen turns into a list with the operators decomposition
+        when only a single layer of expansion is necessary."""
+        op = NoMatOp("a")
+        casted_to_list = list(_operator_decomposition_gen(op))
+        assert len(casted_to_list) == 2
+        assert qml.equal(casted_to_list[0], qml.PauliX("a"))
+        assert qml.equal(casted_to_list[1], qml.PauliY("a"))
+
+    def test_operator_decomposition_gen_decomposed_operator_ragged_nesting(self):
+        """Test that _operator_decomposition_gen handles a decomposition that requires different depths of decomposition."""
+
+        class RaggedDecompositionOp(Operation):
+            num_wires = 1
+
+            def decomposition(self):
+                return [NoMatOp(self.wires), qml.S(self.wires), qml.adjoint(NoMatOp(self.wires))]
+
+        op = RaggedDecompositionOp("a")
+        final_decomp = list(_operator_decomposition_gen(op))
+        assert len(final_decomp) == 5
+        assert qml.equal(final_decomp[0], qml.PauliX("a"))
+        assert qml.equal(final_decomp[1], qml.PauliY("a"))
+        assert qml.equal(final_decomp[2], qml.S("a"))
+        assert qml.equal(final_decomp[3], qml.adjoint(qml.PauliY("a")))
+        assert qml.equal(final_decomp[4], qml.adjoint(qml.PauliX("a")))
+
+    def test_error_from_unsupported_operation(self):
+        """Test that a device error is raised if the operator cant be decomposed and doesn't have a matrix."""
+        op = NoMatNoDecompOp("a")
+        with pytest.raises(DeviceError, match=r"Operator NoMatNoDecompOp"):
+            tuple(_operator_decomposition_gen(op))
+
+
+class TestExpandFnValidation:
+    """Unit tests for helper functions in qml.devices.qubit.preprocess"""
 
-    def test_check_validity_invalid_op(self):
-        """Test that check_validity throws an error when an operation is invalid."""
-        tape = QuantumScript(ops=[qml.Barrier(0)], measurements=[qml.expval(qml.Hadamard(0))])
-        with pytest.raises(DeviceError, match="Gate Barrier not supported on Python Device"):
-            check_validity(tape)
+    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="Operator NoMatNoDecompOp"):
+            expand_fn(tape)
 
-    def test_check_validity_invalid_observable(self):
-        """Test that check_validity throws an error when an observable is invalid."""
+    def test_expand_fn_invalid_observable(self):
+        """Test that expand_fn throws an error when an observable is invalid."""
         tape = QuantumScript(
             ops=[qml.PauliX(0)], measurements=[qml.expval(qml.GellMann(wires=0, index=1))]
         )
-        with pytest.raises(DeviceError, match="Observable GellMann not supported on Python Device"):
-            check_validity(tape)
+        with pytest.raises(DeviceError, match=r"Observable GellMann1"):
+            expand_fn(tape)
 
-    def test_check_validity_invalid_tensor_observable(self):
-        """Test that check_validity throws an error when a tensor includes invalid obserables"""
+    def test_expand_fn_invalid_tensor_observable(self):
+        """Test that expand_fn throws an error when a tensor includes invalid obserables"""
         tape = QuantumScript(
             ops=[qml.PauliX(0), qml.PauliY(1)],
             measurements=[
                 qml.expval(qml.GellMann(wires=0, index=1) @ qml.GellMann(wires=1, index=2))
             ],
         )
-        with pytest.raises(DeviceError, match="Observable GellMann not supported on Python Device"):
-            check_validity(tape)
+        with pytest.raises(DeviceError, match="Observable expval"):
+            expand_fn(tape)
+
+    def test_state_prep_only_one(self):
+        """Test that a device error is raised if the script has multiple state prep ops."""
+        qs = QuantumScript(prep=[qml.BasisState([0], wires=0), qml.BasisState([1], wires=1)])
+        with pytest.raises(
+            DeviceError, match=r"DefaultQubit2 accepts at most one state prep operation."
+        ):
+            expand_fn(qs)
 
-    def test_check_validity_passes(self):
-        """Test that check_validity doesn't throw any errors for a valid circuit"""
+    def test_only_state_based_measurements(self):
+        """Test that a device error is raised if a measurement is not a state measurement."""
+        qs = QuantumScript([], [qml.expval(qml.PauliZ(0)), qml.sample()])
+        with pytest.raises(DeviceError, match=r"Measurement process sample"):
+            expand_fn(qs)
+
+    def test_expand_fn_passes(self):
+        """Test that expand_fn doesn't throw any errors for a valid circuit"""
         tape = QuantumScript(
             ops=[qml.PauliX(0), qml.RZ(0.123, wires=0)], measurements=[qml.state()]
         )
-        check_validity(tape)
+        expand_fn(tape)
+
+    def test_infinite_decomposition_loop(self):
+        """Test that a device error is raised if decomposition enters an infinite loop."""
+
+        class InfiniteOp(qml.operation.Operation):
+            num_wires = 1
+
+            def decomposition(self):
+                return [InfiniteOp(self.wires)]
+
+        qs = qml.tape.QuantumScript([InfiniteOp(0)])
+        with pytest.raises(DeviceError, match=r"Reached recursion limit trying to decompose"):
+            expand_fn(qs)
 
+
+class TestExpandFnTransformations:
     def test_expand_fn_expand_unsupported_op(self):
         """Test that expand_fn expands the tape when unsupported operators are present"""
         ops = [qml.Hadamard(0), NoMatOp(1), qml.RZ(0.123, wires=1)]
@@ -156,6 +220,15 @@ class TestHelpers:
         for op, exp in zip(expanded_tape.circuit, ops + measurements):
             assert qml.equal(op, exp)
 
+    def test_expand_fn_non_commuting_measurements(self):
+        """Test that expand function can decompose operations even when non commuting measurements exist in the circuit."""
+
+        qs = QuantumScript([NoMatOp("a")], [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(0))])
+        new_qs = expand_fn(qs)
+        assert new_qs.measurements == qs.measurements
+
+
+class TestBatchTransform:
     def test_batch_transform_no_batching(self):
         """Test that batch_transform does nothing when no batching is required."""
         ops = [qml.Hadamard(0), qml.CNOT([0, 1]), qml.RX(0.123, wires=1)]
@@ -294,5 +367,5 @@ class TestPreprocess:
             QuantumScript(ops=ops, measurements=measurements[1]),
         ]
 
-        with pytest.raises(qml.DeviceError, match="Gate NoMatNoDecompOp not supported"):
+        with pytest.raises(qml.DeviceError, match="Operator NoMatNoDecompOp"):
             _ = preprocess(tapes)