Repository: PennyLaneAI/pennylane
Pull request #4787 (2023-11-03 22:53:55)

[BUGFIX] Qubit reuse for terminal measurements on wires with mid-circuit measurements
**Context:**
When making terminal measurements on wires with mid-circuit measurements, the circuit was not being transformed correctly. This PR fixes that.

**Description of the Change:**
Wires are considered to be reused by `defer_measurements` when a measured wire has operators *and* terminal measurements applied to it after mid-circuit measurements.

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**

[sc-49507]

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4787?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 [(`3b2caab`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/3b2caab494d4c3a64f6d8319a445e9785ef1ae85?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.64% compared to head [(`638ee4d`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4787?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.63%.
> Report is 1 commits behind head on master.

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


```diff
@@            Coverage Diff             @@
##           master    #4787      +/-   ##
==========================================
- Coverage   99.64%   99.63%   -0.02%     
==========================================
  Files         381      381              
  Lines       34302    34048     -254     
==========================================
- Hits        34181    33924     -257     
- Misses        121      124       +3     
```


| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4787?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/transforms/batch\_transform.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4787?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvYmF0Y2hfdHJhbnNmb3JtLnB5) | `96.03% <100.00%> (+0.03%)` | :arrow_up: |
| [pennylane/transforms/compile.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4787?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvY29tcGlsZS5weQ==) | `100.00% <ø> (ø)` | |
| [pennylane/transforms/defer\_measurements.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4787?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvZGVmZXJfbWVhc3VyZW1lbnRzLnB5) | `100.00% <100.00%> (ø)` | |
| [pennylane/transforms/op\_transforms.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4787?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvb3BfdHJhbnNmb3Jtcy5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/transforms/qfunc\_transforms.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4787?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvcWZ1bmNfdHJhbnNmb3Jtcy5weQ==) | `90.74% <100.00%> (+0.74%)` | :arrow_up: |

... and [42 files with indirect coverage changes](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4787/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/4787?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 a933122ff..83627789c 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -48,6 +48,10 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* `qml.defer_measurements` now correctly transforms circuits when terminal measurements include wires
+  used in mid-circuit measurements.
+  [(#4787)](https://github.com/PennyLaneAI/pennylane/pull/4787)
+
 * Jax jit now works with shot vectors.
   [(#4772)](https://github.com/PennyLaneAI/pennylane/pull/4772/)
 
diff --git a/pennylane/transforms/defer_measurements.py b/pennylane/transforms/defer_measurements.py
index 3626a83d7..aebe49ccc 100644
--- a/pennylane/transforms/defer_measurements.py
+++ b/pennylane/transforms/defer_measurements.py
@@ -66,7 +66,7 @@ def _collect_mid_measure_info(tape: QuantumTape):
     any_repeated_measurements = False
     is_postselecting = False
 
-    for op in tape.operations:
+    for op in tape:
         if isinstance(op, MidMeasureMP):
             if op.postselect is not None:
                 is_postselecting = True


Test patch:
diff --git a/tests/transforms/test_defer_measurements.py b/tests/transforms/test_defer_measurements.py
index ec81062f6..c86fa4b33 100644
--- a/tests/transforms/test_defer_measurements.py
+++ b/tests/transforms/test_defer_measurements.py
@@ -1348,6 +1348,25 @@ class TestQubitReuseAndReset:
             for actual, expected in zip(qnode2.qtape.circuit, expected_circuit)
         )
 
+    def test_measurements_add_new_qubits(self):
+        """Test that qubit reuse related logic is applied if a wire with mid-circuit
+        measurements is included in terminal measurements."""
+        tape = qml.tape.QuantumScript(
+            ops=[qml.Hadamard(0), MidMeasureMP(0)], measurements=[qml.density_matrix(wires=[0])]
+        )
+        expected = np.eye(2) / 2
+
+        tapes, _ = qml.defer_measurements(tape)
+
+        dev = qml.device("default.qubit")
+        res = qml.execute(tapes, dev)
+
+        assert np.allclose(res, expected)
+
+        deferred_tape = tapes[0]
+        assert deferred_tape.operations == [qml.Hadamard(0), qml.CNOT([0, 1])]
+        assert deferred_tape.measurements == [qml.density_matrix(wires=[0])]
+
     def test_wire_is_reset(self):
         """Test that a wire is reset to the |0> state without any local phases
         after measurement if reset is requested."""
@@ -1484,7 +1503,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.probs(wires="a")
+        qml.probs(wires="b")
 
     tape = qml.tape.QuantumScript.from_queue(q)
     tapes, _ = qml.defer_measurements(tape)
@@ -1493,7 +1512,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.probs(wires="a"))
+    assert qml.equal(tape[2], qml.probs(wires="b"))
 
 
 def test_custom_wire_labels_fails_with_reset():