Repository: PennyLaneAI/pennylane
Pull request #5501 (2024-04-11 22:03:40)

Operations/measurements on all wires are drawn correctly with mid-circuit measurements
**Context:**
Bug fix for #5499 

**Description of the Change:**
Adding labels to layer strings only extends to wires, not classical bits.

**Benefits:**
Drawer works correctly with mid-circuit measurements

**Possible Drawbacks:**

**Related GitHub Issues:**


mudit2812: [sc-61127]
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5501?dropdown=coverage&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:
> Project coverage is 99.66%. Comparing base [(`1089bb3`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/1089bb3448081f4fde7bc3b4b792f1abda3555ff?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`e6d2414`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5501?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 6 commits behind head on master.


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


```diff
@@            Coverage Diff             @@
##           master    #5501      +/-   ##
==========================================
- Coverage   99.66%   99.66%   -0.01%     
==========================================
  Files         402      404       +2     
  Lines       37635    37556      -79     
==========================================
- Hits        37510    37430      -80     
- Misses        125      126       +1     
```



</details>

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


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index b117f4b3f..d0edbac6b 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -339,6 +339,9 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* Operators applied to all wires are now drawn correctly in a circuit with mid-circuit measurements.
+  [(#5501)](https://github.com/PennyLaneAI/pennylane/pull/5501)
+
 * Fix a bug where certain unary mid-circuit measurement expressions would raise an uncaught error.
   [(#5480)](https://github.com/PennyLaneAI/pennylane/pull/5480)
 
diff --git a/pennylane/drawer/tape_text.py b/pennylane/drawer/tape_text.py
index 8e635ef50..7069e4749 100644
--- a/pennylane/drawer/tape_text.py
+++ b/pennylane/drawer/tape_text.py
@@ -142,7 +142,8 @@ def _add_op(op, layer_str, config):
 
     label = op.label(decimals=config.decimals, cache=config.cache).replace("\n", "")
     if len(op.wires) == 0:  # operation (e.g. barrier, snapshot) across all wires
-        for i, s in enumerate(layer_str):
+        n_wires = len(config.wire_map)
+        for i, s in enumerate(layer_str[:n_wires]):
             layer_str[i] = s + label
     else:
         for w in op.wires:
@@ -225,7 +226,8 @@ def _add_measurement(m, layer_str, config):
         meas_label = m.return_type.value
 
     if len(m.wires) == 0:  # state or probability across all wires
-        for i, s in enumerate(layer_str):
+        n_wires = len(config.wire_map)
+        for i, s in enumerate(layer_str[:n_wires]):
             layer_str[i] = s + meas_label
 
     for w in m.wires:


Test patch:
diff --git a/tests/drawer/test_draw.py b/tests/drawer/test_draw.py
index 1bd7790f5..b571ed4fd 100644
--- a/tests/drawer/test_draw.py
+++ b/tests/drawer/test_draw.py
@@ -284,6 +284,8 @@ class TestLayering:
 class TestMidCircuitMeasurements:
     """Tests for drawing mid-circuit measurements and classical conditions."""
 
+    # pylint: disable=too-many-public-methods
+
     @pytest.mark.parametrize("device_name", ["default.qubit"])
     def test_qnode_mid_circuit_measurement_not_deferred(self, device_name, mocker):
         """Test that a circuit containing mid-circuit measurements is transformed by the drawer
@@ -329,6 +331,51 @@ class TestMidCircuitMeasurements:
 
         assert drawing == expected_drawing
 
+    @pytest.mark.parametrize(
+        "op", [qml.GlobalPhase(0.1), qml.Identity(), qml.Snapshot(), qml.Barrier()]
+    )
+    @pytest.mark.parametrize("decimals", [None, 2])
+    def test_draw_all_wire_ops(self, op, decimals):
+        """Test that operators acting on all wires are drawn correctly"""
+
+        def func():
+            qml.X(0)
+            qml.X(1)
+            m = qml.measure(0)
+            qml.cond(m, qml.X)(0)
+            qml.apply(op)
+            return qml.expval(qml.Z(0))
+
+        # Stripping to remove trailing white-space because length of white-space at the
+        # end of the drawing depends on the length of each individual line
+        drawing = qml.draw(func, decimals=decimals)().strip()
+        label = op.label(decimals=decimals).replace("\n", "")
+        expected_drawing = (
+            f"0: ──X──┤↗├──X──{label}─┤  <Z>\n1: ──X───║───║──{label}─┤     \n         ╚═══╝"
+        )
+
+        assert drawing == expected_drawing
+
+    @pytest.mark.parametrize(
+        "mp, label", [(qml.sample(), "Sample"), (qml.probs(), "Probs"), (qml.counts(), "Counts")]
+    )
+    def test_draw_all_wire_measurements(self, mp, label):
+        """Test that operators acting on all wires are drawn correctly"""
+
+        def func():
+            qml.X(0)
+            qml.X(1)
+            m = qml.measure(0)
+            qml.cond(m, qml.X)(0)
+            return qml.apply(mp)
+
+        # Stripping to remove trailing white-space because length of white-space at the
+        # end of the drawing depends on the length of each individual line
+        drawing = qml.draw(func)().strip()
+        expected_drawing = f"0: ──X──┤↗├──X─┤  {label}\n1: ──X───║───║─┤  {label}\n         ╚═══╝"
+
+        assert drawing == expected_drawing
+
     def test_draw_mid_circuit_measurement_multiple_wires(self):
         """Test that mid-circuit measurements are correctly drawn in circuits
         with multiple wires."""