Repository: PennyLaneAI/pennylane
Pull request #4901 (2023-11-29 20:41:19)

Draw `qml.cond`: Part 2; `qml.map_wires` correctly maps nested tape wires
**Context:**
This PR follows up on adding support for `qml.cond` with `qml.draw` by fully allowing users to draw multiple conditional operators that are conditioned on multiple mid-circuit measurements. While implementing this, I tried to use `qml.map_wires` on the tape being drawn (which was eventually scrapped as it's too messy to use with nested tapes), and found out that nested tapes were not handled by `map_wires` correctly.

An example of what drawings _can_ look like is shown here:
```python
def circuit():
    qml.RX(0.5, 0)
    qml.RX(0.5, 1)

    m0 = qml.measure(0, reset=True, postselect=1)
    m1 = qml.measure(1)
    qml.cond(m0 & m1, qml.MultiControlledX)(wires=[1, 2, 3, 0], control_values="110")
    qml.cond(m1, qml.PauliZ)(2)

    m2 = qml.measure(1)
    qml.CNOT([3, 2])
    qml.cond(m0, qml.ctrl(qml.SWAP, control=[1, 2], control_values=[True, False]))(wires=[0, 3])
    qml.cond(m2 & m1, qml.PauliY)(0)

    qml.Toffoli([3, 1, 0])

    m3 = qml.measure(3, postselect=0)
    qml.cond(m3, qml.RX)(1.23, 3)
    qml.cond(m3 & m1, qml.Hadamard)(2)

    return qml.expval(qml.PauliZ(0) @ qml.PauliZ(2))

print(qml.draw(circuit)())
```
```
0: ──RX(0.50)──┤↗₁│  │0⟩──────╭X────────────╭SWAP──Y─╭X────────────────────┤ ╭<Z@Z>
1: ──RX(0.50)───║─────────┤↗├─├●─────┤↗├────├●─────║─├●────────────────────┤ │     
2: ─────────────║──────────║──├●──Z───║──╭X─├○─────║─│───────────────────H─┤ ╰<Z@Z>
3: ─────────────║──────────║──╰○──║───║──╰●─╰SWAP──║─╰●──┤↗₀├──RX(1.23)──║─┤       
                ╚══════════║═══╬══║═══║══════╝     ║      ║    ║         ║         
                           ╚═══╩══╩═══║════════════╬══════║════║═════════╣         
                                      ╚════════════╝      ║    ║         ║         
                                                          ╚════╩═════════╝         
```

**Description of the Change:**
* Update when fillers are used for drawing bit lines.
* Update unicode character used when two bit lines cross over.
* Removed use of `defer_measurements` in `qml.draw`. Now, `QNode`s won't defer measurements unconditionally, rather only when `expansion_strategy=device` is used, although this is only true for the new device API.
* Update `map_wires` to correctly append mapped nested tapes to the parent tape.

**Benefits:**
* `qml.draw` works fully with mid-circuit measurements!
* `qml.map_wires` works with nested tapes

**Possible Drawbacks:**

**Related GitHub Issues:**


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.
mudit2812: [sc-49288]

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index aafadb5e1..4d285b888 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -46,9 +46,11 @@
 
 <h4>Drawing and statistics for mid-circuit measurements 🎨</h4>
 
-* `qml.draw` now supports drawing mid-circuit measurements and conditional operators.
+* `qml.draw` and `qml.draw_mpl` now support drawing mid-circuit measurements and conditional operators.
   [(#4775)](https://github.com/PennyLaneAI/pennylane/pull/4775)
   [(#4803)](https://github.com/PennyLaneAI/pennylane/pull/4803)
+  [(#4832)](https://github.com/PennyLaneAI/pennylane/pull/4832)
+  [(#4901)](https://github.com/PennyLaneAI/pennylane/pull/4901)
 
 <h4>Catalyst is seamlessly integrated with PennyLane ⚗️</h4>
 
@@ -187,10 +189,6 @@
 * `AmplitudeEmbedding` now also supports batching when used with Tensorflow.
   [(#4818)](https://github.com/PennyLaneAI/pennylane/pull/4818)
 
-* `qml.draw` and `qml.draw_mpl` now support drawing mid-circuit measurements.
-  [(#4775)](https://github.com/PennyLaneAI/pennylane/pull/4775)
-  [(#4832)](https://github.com/PennyLaneAI/pennylane/pull/4832)
-
 * `qml.ArbitraryUnitary` now supports batching.
   [(#4745)](https://github.com/PennyLaneAI/pennylane/pull/4745)
 
@@ -415,6 +413,9 @@
   interface no longer raises an unexpected error.
   [(#4889)](https://github.com/PennyLaneAI/pennylane/pull/4889)
 
+* `qml.map_wires` no longer fails when mapping nested quantum tapes.
+  [(#4901)](https://github.com/PennyLaneAI/pennylane/pull/4901)
+
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):
diff --git a/pennylane/drawer/draw.py b/pennylane/drawer/draw.py
index 4ead7723e..10337de0d 100644
--- a/pennylane/drawer/draw.py
+++ b/pennylane/drawer/draw.py
@@ -266,15 +266,7 @@ def _draw_qnode(
                 tapes = qnode.construct(args, kwargs)
                 if isinstance(qnode.device, qml.devices.Device):
                     program = qnode.transform_program
-                    if any(
-                        isinstance(op, qml.measurements.MidMeasureMP)
-                        for op in qnode.tape.operations
-                    ):
-                        tapes, _ = qml.defer_measurements(qnode.tape, device=qnode.device)
-                    else:
-                        tapes = [qnode.tape]
-
-                    tapes = program(tapes)
+                    tapes = program([qnode.tape])
 
             finally:
                 qnode.expansion_strategy = original_expansion_strategy
diff --git a/pennylane/drawer/tape_text.py b/pennylane/drawer/tape_text.py
index d77b48618..a06b650de 100644
--- a/pennylane/drawer/tape_text.py
+++ b/pennylane/drawer/tape_text.py
@@ -82,7 +82,8 @@ def _add_cond_grouping_symbols(op, layer_str, config):
 
     for b in range(n_wires, max_b):
         if b - n_wires in mapped_bits:
-            layer_str[b] = "═╣"
+            intersection = "╣" if config.cur_layer == config.final_cond_layers[b - n_wires] else "╬"
+            layer_str[b] = "═" + intersection
         else:
             filler = " " if layer_str[b][-1] == " " else "═"
             layer_str[b] = filler + "║"
@@ -459,7 +460,11 @@ def tape_text(
                 if cur_layer_mid_measure is not None:
                     # This condition is needed to pad the filler on the bits under MidMeasureMPs
                     # that are used for conditions correctly
-                    cur_b_filler = b_filler if bit_map[cur_layer_mid_measure] >= b else " "
+                    cur_b_filler = (
+                        b_filler
+                        if bit_map[cur_layer_mid_measure] >= b and i < final_cond_layers[b]
+                        else " "
+                    )
                 else:
                     cur_b_filler = (
                         b_filler if measurement_layers[b] < i < final_cond_layers[b] else " "
diff --git a/pennylane/ops/functions/map_wires.py b/pennylane/ops/functions/map_wires.py
index 0a2e5ed7e..6e890aca8 100644
--- a/pennylane/ops/functions/map_wires.py
+++ b/pennylane/ops/functions/map_wires.py
@@ -109,7 +109,12 @@ def map_wires(
 def _map_wires_transform(
     tape: qml.tape.QuantumTape, wire_map=None, queue=False
 ) -> (Sequence[qml.tape.QuantumTape], Callable):
-    ops = [map_wires(op, wire_map, queue=queue) for op in tape.operations]
+    ops = [
+        map_wires(op, wire_map, queue=queue)
+        if not isinstance(op, QuantumScript)
+        else map_wires(op, wire_map, queue=queue)[0][0]
+        for op in tape.operations
+    ]
     measurements = [map_wires(m, wire_map, queue=queue) for m in tape.measurements]
 
     out = tape.__class__(


Test patch:
diff --git a/tests/drawer/test_draw.py b/tests/drawer/test_draw.py
index 0ffa8c291..d0da838f0 100644
--- a/tests/drawer/test_draw.py
+++ b/tests/drawer/test_draw.py
@@ -285,7 +285,7 @@ class TestMidCircuitMeasurements:
     """Tests for drawing mid-circuit measurements and classical conditions."""
 
     @pytest.mark.parametrize("device_name", ["default.qubit"])
-    def test_mid_circuit_measurement_device_api(self, device_name, mocker):
+    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
         to use deferred measurements if the device uses the new device API."""
         dev = qml.device(device_name)
@@ -299,8 +299,10 @@ class TestMidCircuitMeasurements:
         draw_qnode = qml.draw(circ)
         spy = mocker.spy(qml.defer_measurements, "_transform")
 
-        _ = draw_qnode()
-        spy.assert_called_once()
+        drawing = draw_qnode()
+        spy.assert_not_called()
+
+        assert drawing == "0: ──X──┤↗├─┤  Probs"
 
     @pytest.mark.parametrize(
         "postselect, reset, mid_measure_label",
@@ -444,6 +446,74 @@ class TestMidCircuitMeasurements:
 
         assert drawing == expected_drawing
 
+    def test_multi_meas_single_cond(self):
+        """Test that a circuit is drawn correctly when a single conditioned operation relies on
+        multiple mid-circuit measurements."""
+
+        def circ():
+            qml.RX(0.5, 0)
+            qml.RX(0.5, 1)
+            m0 = qml.measure(0)
+            m1 = qml.measure(1)
+            qml.CNOT([1, 2])
+            m2 = qml.measure(2)
+            qml.cond(m0 & m1 & m2, qml.RZ)(1.23, 1)
+            return qml.expval(qml.PauliZ(0))
+
+        drawing = qml.draw(circ)()
+        expected_drawing = (
+            "0: ──RX(0.50)──┤↗├────────────────────────┤  <Z>\n"
+            "1: ──RX(0.50)───║───┤↗├─╭●───────RZ(1.23)─┤     \n"
+            "2: ─────────────║────║──╰X──┤↗├──║────────┤     \n"
+            "                ╚════║═══════║═══╣              \n"
+            "                     ╚═══════║═══╣              \n"
+            "                             ╚═══╝              "
+        )
+
+        assert drawing == expected_drawing
+
+    def test_multi_meas_multi_cond(self):
+        """Test that a circuit is drawn correctly when there are multiple conditional operations
+        that rely on single or multiple mid-circuit measurements."""
+
+        def circ():
+            qml.RX(0.5, 0)
+            qml.RX(0.5, 1)
+
+            m0 = qml.measure(0, reset=True, postselect=1)
+            m1 = qml.measure(1)
+            qml.cond(m0 & m1, qml.MultiControlledX)(wires=[1, 2, 3, 0], control_values="110")
+            qml.cond(m1, qml.PauliZ)(2)
+
+            m2 = qml.measure(1)
+            qml.CNOT([3, 2])
+            qml.cond(m0, qml.ctrl(qml.SWAP, control=[1, 2], control_values=[True, False]))(
+                wires=[0, 3]
+            )
+            qml.cond(m2 & m1, qml.PauliY)(0)
+
+            qml.Toffoli([3, 1, 0])
+
+            m3 = qml.measure(3, postselect=0)
+            qml.cond(m3, qml.RX)(1.23, 3)
+            qml.cond(m3 & m1, qml.Hadamard)(2)
+
+            return qml.expval(qml.PauliZ(0) @ qml.PauliZ(2))
+
+        drawing = qml.draw(circ)()
+        expected_drawing = (
+            "0: ──RX(0.50)──┤↗₁│  │0⟩──────╭X────────────╭SWAP──Y─╭X────────────────────┤ ╭<Z@Z>\n"
+            "1: ──RX(0.50)───║─────────┤↗├─├●─────┤↗├────├●─────║─├●────────────────────┤ │     \n"
+            "2: ─────────────║──────────║──├●──Z───║──╭X─├○─────║─│───────────────────H─┤ ╰<Z@Z>\n"
+            "3: ─────────────║──────────║──╰○──║───║──╰●─╰SWAP──║─╰●──┤↗₀├──RX(1.23)──║─┤       \n"
+            "                ╚══════════║═══╬══║═══║══════╝     ║      ║    ║         ║         \n"
+            "                           ╚═══╩══╩═══║════════════╬══════║════║═════════╣         \n"
+            "                                      ╚════════════╝      ║    ║         ║         \n"
+            "                                                          ╚════╩═════════╝         "
+        )
+
+        assert drawing == expected_drawing
+
     def test_single_meas_multi_cond_split_lines(self):
         """Test that a circuit is when multiple lines are needed and the measurement
         and condition are on different lines."""
@@ -493,8 +563,8 @@ class TestMidCircuitMeasurements:
         assert drawing == expected_drawing
 
     def test_single_meas_multi_cond_first_line(self):
-        """Test that a circuit is when multiple lines are needed and the measurement
-        and condition are on the same lines in the first line."""
+        """Test that a circuit is drawn correctly when multiple lines are needed and the
+        measurement and condition are on the same lines in the first line."""
 
         def circ():
             m0 = qml.measure(0)
@@ -516,6 +586,88 @@ class TestMidCircuitMeasurements:
 
         assert drawing == expected_drawing
 
+    def test_multi_meas_single_cond_split_lines(self):
+        """Test that a circuit is drawn correctly when multiple lines are needed and the
+        measurements and condition are split between the lines."""
+
+        def circ():
+            qml.RX(0.5, 0)
+            qml.RX(0.5, 1)
+            m0 = qml.measure(0)
+            m1 = qml.measure(1)
+            qml.CNOT([1, 2])
+            m2 = qml.measure(2)
+            qml.cond(m0 & m1 & m2, qml.RZ)(1.23, 1)
+            return qml.expval(qml.PauliZ(0))
+
+        drawing = qml.draw(circ, max_length=30)()
+        expected_drawing = (
+            "0: ──RX(0.50)──┤↗├────────\n"
+            "1: ──RX(0.50)───║───┤↗├─╭●\n"
+            "2: ─────────────║────║──╰X\n"
+            "                ╚════║════\n"
+            "                     ╚════\n"
+            "                          \n\n"
+            "─────────────────┤  <Z>\n"
+            "────────RZ(1.23)─┤     \n"
+            "───┤↗├──║────────┤     \n"
+            "════║═══╣              \n"
+            "════║═══╣              \n"
+            "    ╚═══╝              "
+        )
+
+        assert drawing == expected_drawing
+
+    def test_multi_meas_multi_cond_split_lines(self):
+        """Test that a circuit is drawn correctly when multiple lines are needed and the
+        measurements and conditions are split between the lines."""
+
+        def circ():
+            qml.RX(0.5, 0)
+            qml.RX(0.5, 1)
+
+            m0 = qml.measure(0, reset=True, postselect=1)
+            m1 = qml.measure(1)
+            qml.cond(m0 & m1, qml.MultiControlledX)(wires=[1, 2, 3, 0], control_values="110")
+            qml.cond(m1, qml.PauliZ)(2)
+
+            m2 = qml.measure(1)
+            qml.CNOT([3, 2])
+            qml.cond(m0, qml.ctrl(qml.SWAP, control=[1, 2], control_values=[True, False]))(
+                wires=[0, 3]
+            )
+            qml.cond(m2 & m1, qml.PauliY)(0)
+
+            qml.Toffoli([3, 1, 0])
+
+            m3 = qml.measure(3, postselect=0)
+            qml.cond(m3, qml.RX)(1.23, 3)
+            qml.cond(m3 & m1, qml.Hadamard)(2)
+
+            return qml.expval(qml.PauliZ(0) @ qml.PauliZ(2))
+
+        drawing = qml.draw(circ, max_length=60)()
+        expected_drawing = (
+            "0: ──RX(0.50)──┤↗₁│  │0⟩──────╭X────────────╭SWAP──Y─╭X\n"
+            "1: ──RX(0.50)───║─────────┤↗├─├●─────┤↗├────├●─────║─├●\n"
+            "2: ─────────────║──────────║──├●──Z───║──╭X─├○─────║─│─\n"
+            "3: ─────────────║──────────║──╰○──║───║──╰●─╰SWAP──║─╰●\n"
+            "                ╚══════════║═══╬══║═══║══════╝     ║   \n"
+            "                           ╚═══╩══╩═══║════════════╬═══\n"
+            "                                      ╚════════════╝   \n"
+            "                                                       \n\n"
+            "─────────────────────┤ ╭<Z@Z>\n"
+            "─────────────────────┤ │     \n"
+            "───────────────────H─┤ ╰<Z@Z>\n"
+            "───┤↗₀├──RX(1.23)──║─┤       \n"
+            "    ║    ║         ║         \n"
+            "════║════║═════════╣         \n"
+            "    ║    ║         ║         \n"
+            "    ╚════╩═════════╝         "
+        )
+
+        assert drawing == expected_drawing
+
 
 @pytest.mark.parametrize(
     "transform",
diff --git a/tests/drawer/test_tape_text.py b/tests/drawer/test_tape_text.py
index 671300d12..8744ea16d 100644
--- a/tests/drawer/test_tape_text.py
+++ b/tests/drawer/test_tape_text.py
@@ -106,7 +106,7 @@ class TestHelperFunctions:  # pylint: disable=too-many-arguments
         )
 
     @pytest.mark.parametrize(
-        "cond_op, args, kwargs, out, bit_map, mv",
+        "cond_op, args, kwargs, out, bit_map, mv, cur_layer",
         [
             (
                 qml.PauliX,
@@ -115,6 +115,7 @@ class TestHelperFunctions:  # pylint: disable=too-many-arguments
                 ["─", "─║", "─║", "─║", "═╩"],
                 cond_bit_map_1,
                 default_measurement_value_1,
+                1,
             ),
             (
                 qml.MultiRZ,
@@ -123,6 +124,7 @@ class TestHelperFunctions:  # pylint: disable=too-many-arguments
                 ["─", "─", "─║", "─║", "═╩"],
                 cond_bit_map_1,
                 default_measurement_value_1,
+                1,
             ),
             (
                 qml.Toffoli,
@@ -131,18 +133,29 @@ class TestHelperFunctions:  # pylint: disable=too-many-arguments
                 ["─", "─", "─", "─║", " ║", "═╝"],
                 cond_bit_map_2,
                 default_measurement_value_2,
+                1,
             ),
             (
                 qml.Toffoli,
                 [],
                 {"wires": [0, 1, 2]},
-                ["─", "─", "─", "─║", "═╣", "═╝"],
+                ["─", "─", "─", "─║", "═╬", "═╝"],
                 cond_bit_map_2,
                 default_measurement_value_1 & default_measurement_value_2,
+                1,
+            ),
+            (
+                qml.Toffoli,
+                [],
+                {"wires": [0, 1, 2]},
+                ["─", "─", "─", "─║", "═╣", "═╩"],
+                cond_bit_map_2,
+                default_measurement_value_1 & default_measurement_value_2,
+                0,
             ),
         ],
     )
-    def test_add_cond_grouping_symbols(self, cond_op, bit_map, mv, args, kwargs, out):
+    def test_add_cond_grouping_symbols(self, cond_op, bit_map, mv, cur_layer, args, kwargs, out):
         """Test private _add_grouping_symbols function renders as expected for Conditionals."""
         op = get_conditional_op(mv, cond_op, *args, **kwargs)
         layer_str = ["─", "─", "─", ""] + [" "] * len(bit_map)
@@ -151,7 +164,10 @@ class TestHelperFunctions:  # pylint: disable=too-many-arguments
             op,
             layer_str,
             _Config(
-                wire_map=default_wire_map, bit_map=bit_map, cur_layer=1, final_cond_layers=[0, 1]
+                wire_map=default_wire_map,
+                bit_map=bit_map,
+                cur_layer=cur_layer,
+                final_cond_layers=[0, 1],
             ),
         )
 
diff --git a/tests/ops/functions/test_map_wires.py b/tests/ops/functions/test_map_wires.py
index 60aafa190..e5f25fe86 100644
--- a/tests/ops/functions/test_map_wires.py
+++ b/tests/ops/functions/test_map_wires.py
@@ -139,6 +139,30 @@ class TestMapWiresTapes:
         assert m_obs.wires == Wires(wire_map[1])
         assert qml.math.allclose(dev.execute(tape), dev.execute(m_tape), atol=0.05)
 
+    def test_map_wires_nested_tape(self):
+        """Tes that the tape wires are correct when it has nested tapes."""
+
+        tape = QuantumScript(
+            [
+                qml.PauliX(0),
+                qml.PauliZ(1),
+                qml.PauliX(3),
+                QuantumScript([qml.PauliY(0), qml.Hadamard(1), qml.PauliY(3)]),
+            ]
+        )
+
+        [m_tape], _ = qml.map_wires(tape, wire_map=wire_map)
+
+        m_ops = m_tape.operations
+        assert len(m_ops) == 4
+        assert len(m_tape.measurements) == 0
+        assert m_ops[:3] == [qml.PauliX(4), qml.PauliZ(3), qml.PauliX(1)]
+
+        nested_m_tape = m_ops[3]
+        assert isinstance(nested_m_tape, QuantumScript)
+        assert nested_m_tape.operations == [qml.PauliY(4), qml.Hadamard(3), qml.PauliY(1)]
+        assert len(nested_m_tape.measurements) == 0
+
 
 class TestMapWiresQNodes:
     """Tests for the qml.map_wires method used with qnodes."""
diff --git a/tests/test_qnode.py b/tests/test_qnode.py
index 08ed2b79a..f645e31ea 100644
--- a/tests/test_qnode.py
+++ b/tests/test_qnode.py
@@ -921,23 +921,6 @@ class TestIntegration:
         assert np.allclose(r2[1], mv_res(first_par))
         assert spy.call_count == call_count  # once for each preprocessing
 
-    def test_drawing_has_deferred_measurements(self):
-        """Test that `qml.draw` with qnodes uses defer_measurements
-        to draw circuits with mid-circuit measurements."""
-        dev = qml.device("default.qubit", wires=2)
-
-        @qml.qnode(dev)
-        def circuit(x):
-            qml.RX(x, wires=0)
-            m = qml.measure(0)
-            qml.cond(m, qml.PauliX)(wires=1)
-            return qml.expval(qml.PauliZ(wires=1))
-
-        res = qml.draw(circuit)("x")
-        expected = "0: ──RX(x)─╭●─┤     \n1: ────────╰X─┤  <Z>"
-
-        assert res == expected
-
     @pytest.mark.parametrize("basis_state", [[1, 0], [0, 1]])
     def test_sampling_with_mcm(self, basis_state, mocker):
         """Tests that a QNode with qml.sample and mid-circuit measurements