Repository: PennyLaneAI/pennylane
Pull request #4737 (2023-10-25 21:13:33)

run qnode transforms on new device with draw_mpl as well
**Context:**
This is already being done in the text drawer, I'm just copying it over for consistency.

**Description of the Change:**
If expansion_strategy is _not_ "device" but we're drawing a QNode with a new device, run the QNode transform program on the tape before drawing it. This is for draw_mpl, it's already present for `draw`

**Benefits:**
Parity between the two draw tools

**Possible Drawbacks:**
N/A

**Related GitHub Issues:**
done in #4595 for the text drawer

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/4737?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: `1 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`43492b0`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/43492b0a1698f1beeaa8960b793913a5e58bf3f0?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.64% compared to head [(`7cf496e`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4737?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.63%.

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


```diff
@@                    Coverage Diff                     @@
##           draw-with-device-wires    #4737      +/-   ##
==========================================================
- Coverage                   99.64%   99.63%   -0.01%     
==========================================================
  Files                         380      380              
  Lines                       33980    33987       +7     
==========================================================
+ Hits                        33858    33864       +6     
- Misses                        122      123       +1     
```


| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4737?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/drawer/draw.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4737?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RyYXdlci9kcmF3LnB5) | `98.92% <87.50%> (-1.08%)` | :arrow_down: |


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4737?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/pennylane/drawer/draw.py b/pennylane/drawer/draw.py
index d5788a5aa..5b3432233 100644
--- a/pennylane/drawer/draw.py
+++ b/pennylane/drawer/draw.py
@@ -577,11 +577,23 @@ def _draw_mpl_qnode(
             try:
                 qnode.expansion_strategy = expansion_strategy or original_expansion_strategy
                 qnode.construct(args, kwargs_qnode)
+                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)
+                    tape = tapes[0]
+                else:
+                    tape = qnode.tape
             finally:
                 qnode.expansion_strategy = original_expansion_strategy
 
-            tape = qnode.tape
-
         _wire_order = wire_order or qnode.device.wires or tape.wires
 
         return tape_mpl(


Test patch:
diff --git a/tests/drawer/test_draw_mpl.py b/tests/drawer/test_draw_mpl.py
index b5551a2bb..d6ec97bf4 100644
--- a/tests/drawer/test_draw_mpl.py
+++ b/tests/drawer/test_draw_mpl.py
@@ -339,3 +339,38 @@ def test_draw_mpl_with_qfunc_warns_with_expansion_strategy():
 
     with pytest.warns(UserWarning, match="the expansion_strategy argument is ignored"):
         _ = qml.draw_mpl(qfunc, expansion_strategy="gradient")
+
+
+def test_mid_circuit_measurement_device_api(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."""
+
+    @qml.qnode(qml.device("default.qubit"))
+    def circ():
+        qml.PauliX(0)
+        qml.measure(0)
+        return qml.probs(wires=0)
+
+    draw_qnode = qml.draw_mpl(circ)
+    spy = mocker.spy(qml.defer_measurements, "_transform")
+
+    _ = draw_qnode()
+    spy.assert_called_once()
+
+
+def test_qnode_transform_program(mocker):
+    """Test that qnode transforms are applied before drawing a circuit."""
+
+    @qml.compile
+    @qml.qnode(qml.device("default.qubit"))
+    def circuit():
+        qml.RX(1.1, 0)
+        qml.RX(2.2, 0)
+        return qml.state()
+
+    draw_qnode = qml.draw_mpl(circuit, decimals=2)
+    qnode_transform = circuit.transform_program[0]
+    spy = mocker.spy(qnode_transform, "_transform")
+
+    _ = draw_qnode()
+    spy.assert_called_once()