Fix qml.draw/_mpl to work with deferred measurements circuits that use MCMs in measurements

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)
+
 * Cast the keys of the `CountsMP` measurements returned `dynamic_one_shot` to the type produced by `MeasurementValue.concretize`.
   [(#5587)](https://github.com/PennyLaneAI/pennylane/pull/5587)
 

pennylane/drawer/tape_mpl.py

@@ -27,7 +27,13 @@
 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)
 

pennylane/drawer/tape_text.py

@@ -31,7 +31,13 @@
 )
 
 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", [])

pennylane/drawer/utils.py

@@ -16,6 +16,7 @@
 """
 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 @@
                         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

tests/transforms/test_defer_measurements.py

@@ -1580,6 +1580,56 @@ def qfunc():
         )
         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)
+
+        dev = qml.device("default.qubit", wires=4)
+
+        transformed_qfunc = qml.transforms.defer_measurements(qfunc)
+        transformed_qnode = qml.QNode(transformed_qfunc, dev)
+
+        spaces = " " * len(label)
+        expected = (
+            f"0: ─╭●─╭X───────────┤  {spaces}\n"
+            f"1: ─│──│──╭RY(0.31)─┤  {spaces}\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)
+
+        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."""