Repository: PennyLaneAI/pennylane
Pull request #4701 (2023-10-20 17:46:48)

`defer_measurements` raises errors for unsupported measurements
**Context:**
Measurements that cover all wires can give users unexpected outputs when using mid-circuit measurements. So, `defer_measurements` now explicitly removes support for such measurements by raising errors.

**Description of the Change:**
* Raise an error if `qml.state()` is used.
* Raise an error if `qml.probs()`, `qml.sample()`, or `qml.counts()` is used if wires or an observable isn't specified.
* Add warning to docs

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**
#4363 


mudit2812: > I know there's a lot of debate over this, but I'm fine with this decision and its better than people being confused about what their answers actually mean.

For reference, @trbromley and I had a discussion offline and the decision was to go in this direction in the end.
Alex-Preciado: @mudit2812 Thank you for validating with @trbromley and for posting the update here.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4701?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`a7114a9`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/a7114a9b336e06ca2bb736e1a069d06cb6af3951?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.64% compared to head [(`a4cf441`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4701?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.64%.
> Report is 2 commits behind head on master.

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


```diff
@@            Coverage Diff             @@
##           master    #4701      +/-   ##
==========================================
- Coverage   99.64%   99.64%   -0.01%     
==========================================
  Files         379      380       +1     
  Lines       34126    33983     -143     
==========================================
- Hits        34005    33861     -144     
- Misses        121      122       +1     
```


| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4701?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/\_device.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4701?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19kZXZpY2UucHk=) | `99.41% <100.00%> (+<0.01%)` | :arrow_up: |
| [pennylane/devices/default\_qubit.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4701?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdC5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/devices/qubit/sampling.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4701?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvcXViaXQvc2FtcGxpbmcucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/devices/qubit/simulate.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4701?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvcXViaXQvc2ltdWxhdGUucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/drawer/draw.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4701?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RyYXdlci9kcmF3LnB5) | `100.00% <100.00%> (ø)` | |
| [pennylane/math/single\_dispatch.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4701?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvc2luZ2xlX2Rpc3BhdGNoLnB5) | `99.48% <100.00%> (+<0.01%)` | :arrow_up: |
| [pennylane/measurements/mid\_measure.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4701?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmVtZW50cy9taWRfbWVhc3VyZS5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/operation.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4701?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wZXJhdGlvbi5weQ==) | `97.27% <100.00%> (-0.02%)` | :arrow_down: |
| [pennylane/ops/op\_math/composite.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4701?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9vcF9tYXRoL2NvbXBvc2l0ZS5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/ops/op\_math/controlled.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4701?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9vcF9tYXRoL2NvbnRyb2xsZWQucHk=) | `99.63% <ø> (-0.01%)` | :arrow_down: |
| ... and [7 more](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4701?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | |

... and [39 files with indirect coverage changes](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4701/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4701?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 812e35faa..71c11cbab 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -382,9 +382,12 @@
   decomposition.
   [(#4675)](https://github.com/PennyLaneAI/pennylane/pull/4675)
 
-
 <h3>Breaking changes 💔</h3>
 
+* ``qml.defer_measurements`` now raises an error if a transformed circuit measures ``qml.probs``,
+  ``qml.sample``, or ``qml.counts`` without any wires or obsrvable, or if it measures ``qml.state``.
+  [(#4701)](https://github.com/PennyLaneAI/pennylane/pull/4701)
+
 * The device test suite now converts device kwargs to integers or floats if they can be converted to integers or floats.
   [(#4640)](https://github.com/PennyLaneAI/pennylane/pull/4640)
 
diff --git a/pennylane/devices/default_qubit.py b/pennylane/devices/default_qubit.py
index 52bcd9047..c1a21cbb4 100644
--- a/pennylane/devices/default_qubit.py
+++ b/pennylane/devices/default_qubit.py
@@ -387,8 +387,8 @@ class DefaultQubit(Device):
         config = self._setup_execution_config(execution_config)
         transform_program = TransformProgram()
 
-        transform_program.add_transform(qml.defer_measurements, device=self)
         transform_program.add_transform(validate_device_wires, self.wires, name=self.name)
+        transform_program.add_transform(qml.defer_measurements, device=self)
         transform_program.add_transform(
             decompose, stopping_condition=stopping_condition, name=self.name
         )
diff --git a/pennylane/transforms/defer_measurements.py b/pennylane/transforms/defer_measurements.py
index 1acf9ec29..c39bb0819 100644
--- a/pennylane/transforms/defer_measurements.py
+++ b/pennylane/transforms/defer_measurements.py
@@ -14,7 +14,7 @@
 """Code for the tape transform implementing the deferred measurement principle."""
 from typing import Sequence, Callable
 import pennylane as qml
-from pennylane.measurements import MidMeasureMP
+from pennylane.measurements import MidMeasureMP, ProbabilityMP, SampleMP, CountsMP
 from pennylane.ops.op_math import ctrl
 
 from pennylane.tape import QuantumTape
@@ -23,7 +23,66 @@ from pennylane.transforms.core import transform
 from pennylane.wires import Wires
 from pennylane.queuing import QueuingManager
 
-# pylint: disable=too-many-branches, too-many-statements
+# pylint: disable=too-many-branches, protected-access
+
+
+def _check_tape_validity(tape: QuantumTape):
+    """Helper function to check that the tape is valid."""
+    cv_types = (qml.operation.CVOperation, qml.operation.CVObservable)
+    ops_cv = any(isinstance(op, cv_types) and op.name != "Identity" for op in tape.operations)
+    obs_cv = any(
+        isinstance(getattr(op, "obs", None), cv_types)
+        and not isinstance(getattr(op, "obs", None), qml.Identity)
+        for op in tape.measurements
+    )
+    if ops_cv or obs_cv:
+        raise ValueError("Continuous variable operations and observables are not supported.")
+
+    for mp in tape.measurements:
+        if isinstance(mp, (CountsMP, ProbabilityMP, SampleMP)) and not (
+            mp.obs or mp._wires or mp.mv
+        ):
+            raise ValueError(
+                f"Cannot use {mp.__class__.__name__} as a measurement without specifying wires "
+                "when using qml.defer_measurements. Deferred measurements can occur "
+                "automatically when using mid-circuit measurements on a device that does not "
+                "support them."
+            )
+
+        if mp.__class__.__name__ == "StateMP":
+            raise ValueError(
+                "Cannot use StateMP as a measurement when using qml.defer_measurements. "
+                "Deferred measurements can occur automatically when using mid-circuit "
+                "measurements on a device that does not support them."
+            )
+
+
+def _collect_mid_measure_info(tape: QuantumTape):
+    """Helper function to collect information related to mid-circuit measurements in the tape."""
+
+    # Find wires that are reused after measurement
+    measured_wires = []
+    reused_measurement_wires = set()
+    any_repeated_measurements = False
+    is_postselecting = False
+
+    for op in tape.operations:
+        if isinstance(op, MidMeasureMP):
+            if op.postselect is not None:
+                is_postselecting = True
+            if op.reset:
+                reused_measurement_wires.add(op.wires[0])
+
+            if op.wires[0] in measured_wires:
+                any_repeated_measurements = True
+            measured_wires.append(op.wires[0])
+
+        else:
+            reused_measurement_wires = reused_measurement_wires.union(
+                set(measured_wires).intersection(op.wires.toset())
+            )
+
+    return measured_wires, reused_measurement_wires, any_repeated_measurements, is_postselecting
 
 
 def null_postprocessing(results):
@@ -68,12 +127,18 @@ def defer_measurements(tape: QuantumTape, **kwargs) -> (Sequence[QuantumTape], C
 
     .. note::
 
-        When applying the transform on a quantum function that returns
-        :func:`~pennylane.state` as the terminal measurement or contains the
+        When applying the transform on a quantum function that contains the
         :class:`~.Snapshot` instruction, state information corresponding to
         simulating the transformed circuit will be obtained. No
         post-measurement states are considered.
 
+    .. warning::
+
+        :func:`~.pennylane.state` is not supported with the ``defer_measurements`` transform.
+        Additionally, :func:`~.pennylane.probs`, :func:`~.pennylane.sample` and
+        :func:`~.pennylane.counts` can only be used with ``defer_measurements`` if wires
+        or an observable are explicitly specified.
+
     Args:
         tape (.QuantumTape): a quantum tape
 
@@ -136,45 +201,25 @@ def defer_measurements(tape: QuantumTape, **kwargs) -> (Sequence[QuantumTape], C
     >>> func(*pars)
     tensor([0.76960924, 0.13204407, 0.08394415, 0.01440254], requires_grad=True)
     """
-    # pylint: disable=protected-access
+
     if not any(isinstance(o, MidMeasureMP) for o in tape.operations):
         return (tape,), null_postprocessing
 
-    cv_types = (qml.operation.CVOperation, qml.operation.CVObservable)
-    ops_cv = any(isinstance(op, cv_types) and op.name != "Identity" for op in tape.operations)
-    obs_cv = any(
-        isinstance(getattr(op, "obs", None), cv_types)
-        and not isinstance(getattr(op, "obs", None), qml.Identity)
-        for op in tape.measurements
-    )
-    if ops_cv or obs_cv:
-        raise ValueError("Continuous variable operations and observables are not supported.")
+    _check_tape_validity(tape)
+
+    device = kwargs.get("device", None)
 
     device = kwargs.get("device", None)
 
     new_operations = []
 
     # Find wires that are reused after measurement
-    measured_wires = []
-    reused_measurement_wires = set()
-    repeated_measurement_wire = False
-    is_postselecting = False
-
-    for op in tape.operations:
-        if isinstance(op, MidMeasureMP):
-            if op.postselect is not None:
-                is_postselecting = True
-            if op.reset:
-                reused_measurement_wires.add(op.wires[0])
-
-            if op.wires[0] in measured_wires:
-                repeated_measurement_wire = True
-            measured_wires.append(op.wires[0])
-
-        else:
-            reused_measurement_wires = reused_measurement_wires.union(
-                set(measured_wires).intersection(op.wires.toset())
-            )
+    (
+        measured_wires,
+        reused_measurement_wires,
+        any_repeated_measurements,
+        is_postselecting,
+    ) = _collect_mid_measure_info(tape)
 
     if is_postselecting and device is not None and not isinstance(device, qml.devices.DefaultQubit):
         raise ValueError(f"Postselection is not supported on the {device} device.")
@@ -183,7 +228,7 @@ def defer_measurements(tape: QuantumTape, **kwargs) -> (Sequence[QuantumTape], C
     # classically controlled operations
     control_wires = {}
     cur_wire = (
-        max(tape.wires) + 1 if reused_measurement_wires or repeated_measurement_wire else None
+        max(tape.wires) + 1 if reused_measurement_wires or any_repeated_measurements else None
     )
 
     for op in tape.operations:
@@ -257,7 +302,7 @@ def _add_control_gate(op, control_wires):
     control = [control_wires[m.id] for m in op.meas_val.measurements]
     new_ops = []
 
-    for branch, value in op.meas_val._items():  # pylint: disable=protected-access
+    for branch, value in op.meas_val._items():
         if value:
             qscript = qml.tape.make_qscript(
                 ctrl(


Test patch:
diff --git a/tests/devices/default_qubit/test_default_qubit.py b/tests/devices/default_qubit/test_default_qubit.py
index c56966d49..257d57089 100644
--- a/tests/devices/default_qubit/test_default_qubit.py
+++ b/tests/devices/default_qubit/test_default_qubit.py
@@ -1650,7 +1650,6 @@ class TestPostselection:
             qml.expval(qml.PauliZ(0)),
             qml.var(qml.PauliZ(0)),
             qml.probs(wires=[0, 1]),
-            qml.state(),
             qml.density_matrix(wires=0),
             qml.purity(0),
             qml.vn_entropy(0),
@@ -1800,7 +1799,6 @@ class TestPostselection:
             (qml.expval(qml.PauliZ(0)), "autograd", True),
             (qml.var(qml.PauliZ(0)), "autograd", True),
             (qml.probs(wires=[0, 1]), "autograd", True),
-            (qml.state(), "autograd", True),
             (qml.density_matrix(wires=0), "autograd", True),
             (qml.purity(0), "numpy", True),
             (qml.vn_entropy(0), "numpy", False),
diff --git a/tests/transforms/test_defer_measurements.py b/tests/transforms/test_defer_measurements.py
index f779c5a94..ec81062f6 100644
--- a/tests/transforms/test_defer_measurements.py
+++ b/tests/transforms/test_defer_measurements.py
@@ -80,6 +80,23 @@ def test_postselection_error_with_wrong_device():
         _ = circ()
 
 
+@pytest.mark.parametrize(
+    "mp, err_msg",
+    [
+        (qml.state(), "Cannot use StateMP as a measurement when"),
+        (qml.probs(), "Cannot use ProbabilityMP as a measurement without"),
+        (qml.sample(), "Cannot use SampleMP as a measurement without"),
+        (qml.counts(), "Cannot use CountsMP as a measurement without"),
+    ],
+)
+def test_unsupported_measurements(mp, err_msg):
+    """Test that using unsupported measurements raises an error."""
+    tape = qml.tape.QuantumScript([MidMeasureMP(0)], [mp])
+
+    with pytest.raises(ValueError, match=err_msg):
+        _, _ = qml.defer_measurements(tape)
+
+
 class TestQNode:
     """Test that the transform integrates well with QNodes."""
 
@@ -1467,7 +1484,7 @@ def test_custom_wire_labels_allowed_without_reset():
         qml.Hadamard("a")
         ma = qml.measure("a", reset=False)
         qml.cond(ma, qml.PauliX)("b")
-        qml.state()
+        qml.probs(wires="a")
 
     tape = qml.tape.QuantumScript.from_queue(q)
     tapes, _ = qml.defer_measurements(tape)
@@ -1476,7 +1493,7 @@ def test_custom_wire_labels_allowed_without_reset():
     assert len(tape) == 3
     assert qml.equal(tape[0], qml.Hadamard("a"))
     assert qml.equal(tape[1], qml.CNOT(["a", "b"]))
-    assert qml.equal(tape[2], qml.state())
+    assert qml.equal(tape[2], qml.probs(wires="a"))
 
 
 def test_custom_wire_labels_fails_with_reset():
@@ -1485,7 +1502,7 @@ def test_custom_wire_labels_fails_with_reset():
         qml.Hadamard("a")
         ma = qml.measure("a", reset=True)
         qml.cond(ma, qml.PauliX)("b")
-        qml.state()
+        qml.probs(wires="a")
 
     tape = qml.tape.QuantumScript.from_queue(q)
     with pytest.raises(TypeError, match="can only concatenate str"):