Repository: PennyLaneAI/pennylane
Pull request #5610 (2024-04-30 21:25:53)

Fix `qml.draw/_mpl` to work with deferred measurements circuits that use MCMs in measurements
**Context:**
`defer_measurements` leaves `MeasurementValue`s in the terminal measurement processes if they are present. This is not friendly with the drawer as the classical wires for those measurement values do not exist when using `defer_measurements`. This PR adds a "hack" fix so that the drawer/MPL drawer work with this scenario.

**Description of the Change:**
Add `transform_deferred_measurements_tape` util function, which replaces any `MeasurementValue`s present in the tape after applying `defer_measurements` with wires. This transform is only applied in `tape_text` and `tape_mpl`.

**Benefits:**
Drawer works better with MCMs

**Possible Drawbacks:**
Hacky fix, technical debt to implement cleaner fix later

**Related GitHub Issues:**
https://github.com/PennyLaneAI/pennylane/issues/5588


github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5610?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 `93.75000%` with `1 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/5610?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/drawer/utils.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5610?src=pr&el=tree&filepath=pennylane%2Fdrawer%2Futils.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RyYXdlci91dGlscy5weQ==) | 91.66% | [1 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5610?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    #5610   +/-   ##
==============================================
  Coverage               ?   99.68%           
==============================================
  Files                  ?      412           
  Lines                  ?    38298           
  Branches               ?        0           
==============================================
  Hits                   ?    38179           
  Misses                 ?      119           
  Partials               ?        0           
```



</details>

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

mudit2812: [sc-62204]

Patch:
diff --git a/doc/releases/changelog-0.36.0.md b/doc/releases/changelog-0.36.0.md
index daa240dae..98049b16b 100644
--- a/doc/releases/changelog-0.36.0.md
+++ b/doc/releases/changelog-0.36.0.md
@@ -551,6 +551,10 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* Fixed a bug where `qml.draw` and `qml.draw_mpl` incorrectly raised errors for circuits collecting statistics on mid-circuit measurements
+  while using `qml.defer_measurements`.
+  [(#5610)](https://github.com/PennyLaneAI/pennylane/pull/5610)
+
 * Using shot vectors with `param_shift(... broadcast=True)` caused a bug. This combination is no longer supported
   and will be added again in the next release.
   [(#5612)](https://github.com/PennyLaneAI/pennylane/pull/5612)
diff --git a/pennylane/drawer/tape_mpl.py b/pennylane/drawer/tape_mpl.py
index b6386cdb3..4d9015366 100644
--- a/pennylane/drawer/tape_mpl.py
+++ b/pennylane/drawer/tape_mpl.py
@@ -27,7 +27,13 @@ from pennylane import ops
 from pennylane.measurements import MidMeasureMP
 from .mpldrawer import MPLDrawer
 from .drawable_layers import drawable_layers
-from .utils import convert_wire_order, unwrap_controls, cwire_connections, default_bit_map
+from .utils import (
+    convert_wire_order,
+    cwire_connections,
+    default_bit_map,
+    transform_deferred_measurements_tape,
+    unwrap_controls,
+)
 from .style import _set_style
 
 has_mpl = True
@@ -216,6 +222,7 @@ def _tape_mpl(tape, wire_order=None, show_all_wires=False, decimals=None, *, fig
     fontsize = kwargs.get("fontsize", None)
 
     wire_map = convert_wire_order(tape, wire_order=wire_order, show_all_wires=show_all_wires)
+    tape = transform_deferred_measurements_tape(tape)
     tape = qml.map_wires(tape, wire_map=wire_map)[0][0]
     bit_map = default_bit_map(tape)
 
diff --git a/pennylane/drawer/tape_text.py b/pennylane/drawer/tape_text.py
index 7069e4749..575448312 100644
--- a/pennylane/drawer/tape_text.py
+++ b/pennylane/drawer/tape_text.py
@@ -31,7 +31,13 @@ from pennylane.measurements import (
 )
 
 from .drawable_layers import drawable_layers
-from .utils import convert_wire_order, unwrap_controls, cwire_connections, default_bit_map
+from .utils import (
+    convert_wire_order,
+    cwire_connections,
+    default_bit_map,
+    transform_deferred_measurements_tape,
+    unwrap_controls,
+)
 
 
 @dataclass
@@ -425,6 +431,7 @@ def tape_text(
         New tape offset:  4
 
     """
+    tape = transform_deferred_measurements_tape(tape)
     cache = cache or {}
     cache.setdefault("tape_offset", 0)
     cache.setdefault("matrices", [])
diff --git a/pennylane/drawer/utils.py b/pennylane/drawer/utils.py
index 5d44feb04..46c33cd34 100644
--- a/pennylane/drawer/utils.py
+++ b/pennylane/drawer/utils.py
@@ -16,6 +16,7 @@ This module contains some useful utility functions for circuit drawing.
 """
 from pennylane.ops import Controlled, Conditional
 from pennylane.measurements import MeasurementProcess, MidMeasureMP, MeasurementValue
+from pennylane.tape import QuantumScript
 
 
 def default_wire_map(tape):
@@ -201,3 +202,22 @@ def cwire_connections(layers, bit_map):
                         connected_layers[cwire].append(layer_idx)
 
     return connected_layers, connected_wires
+
+
+def transform_deferred_measurements_tape(tape):
+    """Helper function to replace MeasurementValues with wires for tapes using
+    deferred measurements."""
+    if not any(isinstance(op, MidMeasureMP) for op in tape.operations) and any(
+        m.mv is not None for m in tape.measurements
+    ):
+        new_measurements = []
+        for m in tape.measurements:
+            if m.mv is not None:
+                new_m = type(m)(wires=m.wires)
+                new_measurements.append(new_m)
+            else:
+                new_measurements.append(m)
+        new_tape = QuantumScript(tape.operations, new_measurements)
+        return new_tape
+
+    return tape


Test patch:
diff --git a/tests/transforms/test_defer_measurements.py b/tests/transforms/test_defer_measurements.py
index e72508b3e..8cb00319f 100644
--- a/tests/transforms/test_defer_measurements.py
+++ b/tests/transforms/test_defer_measurements.py
@@ -1580,6 +1580,57 @@ class TestDrawing:
         )
         assert qml.draw(transformed_qnode)() == expected
 
+    @pytest.mark.parametrize(
+        "mp, label",
+        [
+            (qml.sample, "Sample"),
+            (qml.probs, "Probs"),
+            (qml.var, "Var[None]"),
+            (qml.counts, "Counts"),
+            (qml.expval, "<None>"),
+        ],
+    )
+    def test_drawing_with_mcm_terminal_measure(self, mp, label):
+        """Test that drawing a func works correctly when collecting statistics on
+        mid-circuit measurements."""
+
+        def qfunc():
+            m_0 = qml.measure(0, reset=True)
+            qml.cond(m_0, qml.RY)(0.312, wires=1)
+
+            return mp(op=m_0), qml.expval(qml.Z(1))
+
+        dev = qml.device("default.qubit", wires=4)
+
+        transformed_qfunc = qml.transforms.defer_measurements(qfunc)
+        transformed_qnode = qml.QNode(transformed_qfunc, dev)
+
+        spaces = " " * len(label)
+        expval = "<Z>".ljust(len(label))
+        expected = (
+            f"0: ─╭●─╭X───────────┤  {spaces}\n"
+            f"1: ─│──│──╭RY(0.31)─┤  {expval}\n"
+            f"2: ─╰X─╰●─╰●────────┤  {label}"
+        )
+        assert qml.draw(transformed_qnode)() == expected
+
+    @pytest.mark.parametrize("mp", [qml.sample, qml.probs, qml.var, qml.counts, qml.expval])
+    def test_draw_mpl_with_mcm_terminal_measure(self, mp):
+        """Test that no error is raised when drawing a circuit which collects
+        statistics on mid-circuit measurements"""
+
+        def qfunc():
+            m_0 = qml.measure(0, reset=True)
+            qml.cond(m_0, qml.RY)(0.312, wires=1)
+
+            return mp(op=m_0), qml.expval(qml.Z(1))
+
+        dev = qml.device("default.qubit", wires=4)
+
+        transformed_qfunc = qml.transforms.defer_measurements(qfunc)
+        transformed_qnode = qml.QNode(transformed_qfunc, dev)
+        _ = qml.draw_mpl(transformed_qnode)()
+
 
 def test_custom_wire_labels_allowed_without_reuse():
     """Test that custom wire labels work if no qubits are re-used."""