Repository: PennyLaneAI/pennylane (python)

Issue:
[BUG] Issues with custom devices and transforms
### Expected behavior

Transforms applied to QNodes should work the same, independently of which device was used to create the QNode.

### Actual behavior

Executing the code provided below reveals a number of issues. The output I am getting (for example with PL v0.34.0 - but older version back to 0.30 behave essentially the same) is:
```
<default.qubit device (wires=1) at 0x7fae91eb3f10>
0: ──X─┤  <Z>
-1.0 <class 'pennylane.numpy.tensor.tensor'>
0: ──Y──X─┤  <Z>
1.0 <class 'pennylane.numpy.tensor.tensor'>
MyDevice
Short name: my.device
Package: __main__
Plugin version: 0.1.0
Author: Christian Gogolin
Wires: 1
Shots: 1000
0: ──X─┤  <Z>
[-1.] <class 'pennylane.numpy.tensor.tensor'>
0: ──X─┤  <Z>
[1.] <class 'pennylane.numpy.tensor.tensor'>
```
As you can see, when executing the transformed qnode the device receives the correct gates and produces the correct output (up to an odd mismatch in the shape of the returned tensor).

But when drawing the transformed qnode the transform does not seem to have any effect.
To me this is very mysterious because the device is not even needed for drawing the circuit.

With 0.29.1 my code imports the old batch_transform and the transformed circuit is drawn correctly.

### Additional information

_No response_

### Source code

```shell
from typing import Callable, Sequence

import pennylane as qml
try:
    from pennylane.transforms import transform
except ImportError: # pennylane < 0.31
    from pennylane.transforms import batch_transform as transform

from pennylane.tape import QuantumTape

from pennylane import Device
from pennylane import numpy as np

class MyDevice(Device):
    author = "Christian Gogolin"
    pennylane_requires = ">=0.29.1"
    name = "MyDevice"
    short_name = "my.device"
    version = "0.1.0"

    def __init__(self, wires, *args, **kwargs):
        if wires != 1:
            raise NotImplementedError()
        self._state = None
        super().__init__(wires, *args, **kwargs)

    def apply(self, operation, wires, par):
        self._state = qml.matrix(getattr(qml, operation))(0) @ self._state

    def expval(self, observable, wires, par):
        return float(np.real(np.dot(self._state.conj().T, qml.matrix(qml.PauliZ)(0).dot(self._state))).item())

    @property
    def observables(self):
        return ["PauliZ"]

    @property
    def operations(self):
        return ["PauliX", "PauliY"]

    def reset(self):
        self._state = np.array([1., 0.])

for dev in [qml.device("default.qubit", wires=1), MyDevice(wires=1)]:

    @qml.qnode(dev)
    def qnode():
        qml.PauliX(0)
        return qml.expval(qml.PauliZ(0))

    @transform
    def prepend_y(tape: QuantumTape) -> (Sequence[QuantumTape], Callable):

        try:
            new_tape = type(tape)([qml.PauliY(0)] + tape.operations, tape.measurements, shots=tape.shots)
        except AttributeError: # pennylane < 0.31
            new_tape = type(tape)([qml.PauliY(0)] + tape.operations, tape.measurements)

        def postprocessing(results):
            return results[0]

        return [new_tape], postprocessing

    print(dev)
    print(qml.draw(qnode)())
    res = qnode()
    print(res, type(res))
    print(qml.draw(prepend_y(qnode))())
    res2 = prepend_y(qnode)()
    print(res2, type(res2))
```


### Tracebacks

_No response_

### System information

```shell
Python 3.9.16 | packaged by conda-forge | (main, Feb  1 2023, 21:39:03)
[GCC 11.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import pennylane as qml; qml.about()
Name: PennyLane
Version: 0.34.0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/PennyLaneAI/pennylane
Author:
Author-email:
License: Apache License 2.0
Location: /home/cvjjm/miniforge3/envs/hfak/lib/python3.9/site-packages
Editable project location: /home/cvjjm/src/covqcstack/hfak/pennylane
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: covvqetools, PennyLane-Lightning, pennylane-umz

Platform info:           Linux-5.15.133.1-microsoft-standard-WSL2-x86_64-with-glibc2.31
Python version:          3.9.16
Numpy version:           1.23.0
Scipy version:           1.10.0
Installed devices:
- default.gaussian (PennyLane-0.34.0)
- default.mixed (PennyLane-0.34.0)
- default.qubit (PennyLane-0.34.0)
- default.qubit.autograd (PennyLane-0.34.0)
- default.qubit.jax (PennyLane-0.34.0)
- default.qubit.legacy (PennyLane-0.34.0)
- default.qubit.tf (PennyLane-0.34.0)
- default.qubit.torch (PennyLane-0.34.0)
- default.qutrit (PennyLane-0.34.0)
- null.qubit (PennyLane-0.34.0)
- lightning.qubit (PennyLane-Lightning-0.34.0)
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.

Pull request description:
Apply all user transforms before drawing
Fixes #5275 [sc-57775]

Basically, we were only applying the transforms before drawing when the user provided `expansion_strategy` wasn't `"device"` and the device was from the new device interface.  This PR makes it so that we always apply all the user transforms before drawing.  This is a quick fix until #5139 can get in.

Patch:
diff --git a/doc/releases/changelog-0.35.0.md b/doc/releases/changelog-0.35.0.md
index 12758f0e0..09b800936 100644
--- a/doc/releases/changelog-0.35.0.md
+++ b/doc/releases/changelog-0.35.0.md
@@ -552,6 +552,9 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* `qml.draw` and `qml.draw_mpl` now apply all applied transforms before drawing.
+  [(#5277)](https://github.com/PennyLaneAI/pennylane/pull/5277/)
+
 * Fixes a bug in the matplotlib drawer where the color of `Barrier` did not match the requested style.
   [(#5276)](https://github.com/PennyLaneAI/pennylane/pull/5276)
 
diff --git a/pennylane/drawer/draw.py b/pennylane/drawer/draw.py
index cd6d23bac..41fb7fde9 100644
--- a/pennylane/drawer/draw.py
+++ b/pennylane/drawer/draw.py
@@ -259,16 +259,16 @@ def _draw_qnode(
             expansion_strategy == "device" or getattr(qnode, "expansion_strategy", None) == "device"
         ):
             qnode.construct(args, kwargs)
+            tapes = qnode.transform_program([qnode.tape])[0]
             program, _ = qnode.device.preprocess()
-            tapes = program([qnode.tape])
+            tapes = program(tapes)[0]
         else:
             original_expansion_strategy = getattr(qnode, "expansion_strategy", None)
             try:
                 qnode.expansion_strategy = expansion_strategy or original_expansion_strategy
                 tapes = qnode.construct(args, kwargs)
-                if isinstance(qnode.device, qml.devices.Device):
-                    program = qnode.transform_program
-                    tapes = program([qnode.tape])
+                program = qnode.transform_program
+                tapes = program([qnode.tape])[0]
 
             finally:
                 qnode.expansion_strategy = original_expansion_strategy
@@ -287,7 +287,7 @@ def _draw_qnode(
                     max_length=max_length,
                     cache=cache,
                 )
-                for t in tapes[0]
+                for t in tapes
             ]
             if show_matrices and cache["matrices"]:
                 mat_str = ""
@@ -570,8 +570,9 @@ def _draw_mpl_qnode(
     def wrapper(*args, **kwargs_qnode):
         if expansion_strategy == "device" and isinstance(qnode.device, qml.devices.Device):
             qnode.construct(args, kwargs)
+            tapes, _ = qnode.transform_program([qnode.tape])
             program, _ = qnode.device.preprocess()
-            tapes, _ = program([qnode.tape])
+            tapes, _ = program(tapes)
             tape = tapes[0]
         else:
             original_expansion_strategy = getattr(qnode, "expansion_strategy", None)
@@ -579,11 +580,8 @@ 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
-                    [tape], _ = program([qnode.tape])
-                else:
-                    tape = qnode.tape
+                program = qnode.transform_program
+                [tape], _ = program([qnode.tape])
             finally:
                 qnode.expansion_strategy = original_expansion_strategy
 


Test patch:
diff --git a/tests/drawer/test_draw.py b/tests/drawer/test_draw.py
index 1678f47bb..cb6d59046 100644
--- a/tests/drawer/test_draw.py
+++ b/tests/drawer/test_draw.py
@@ -881,6 +881,29 @@ def test_expansion_strategy(device):
     assert draw(circ, expansion_strategy="device", decimals=None)(0.5) == expected_device
 
 
+@pytest.mark.parametrize(
+    "device",
+    [qml.device("default.qubit.legacy", wires=2), qml.device("default.qubit", wires=2)],
+)
+def test_applied_transforms(device):
+    """Test that any transforms applied to the qnode are included in the output."""
+
+    @qml.transform
+    def just_pauli_x(_):
+        new_tape = qml.tape.QuantumScript([qml.PauliX(0)])
+        return (new_tape,), lambda res: res[0]
+
+    @just_pauli_x
+    @qml.qnode(device)
+    def my_circuit(x):
+        qml.RX(x, wires=0)
+        qml.SWAP(wires=(0, 1))
+        return qml.probs(wires=(0, 1))
+
+    expected = "0: ──X─┤  "
+    assert qml.draw(my_circuit)(1.234) == expected
+
+
 def test_draw_with_qfunc():
     """Test a non-qnode qfunc can be drawn."""
 
diff --git a/tests/drawer/test_draw_mpl.py b/tests/drawer/test_draw_mpl.py
index 6308a2b5f..010b38e23 100644
--- a/tests/drawer/test_draw_mpl.py
+++ b/tests/drawer/test_draw_mpl.py
@@ -398,3 +398,31 @@ def test_draw_mpl_with_control_in_adjoint():
     assert len(ax.lines) == 4  # three wires, one control
     assert len(ax.texts) == 4  # three wire labels, one gate label
     assert ax.texts[-1].get_text() == "X†"
+
+
+@pytest.mark.parametrize(
+    "device",
+    [qml.device("default.qubit.legacy", wires=2), qml.device("default.qubit", wires=2)],
+)
+def test_applied_transforms(device):
+    """Test that any transforms applied to the qnode are included in the output."""
+
+    @qml.transform
+    def just_pauli_x(_):
+        new_tape = qml.tape.QuantumScript([qml.PauliX(0)])
+        return (new_tape,), lambda res: res[0]
+
+    @just_pauli_x
+    @qml.qnode(device)
+    def my_circuit():
+        qml.SWAP(wires=(0, 1))
+        qml.CNOT(wires=(0, 1))
+        return qml.probs(wires=(0, 1))
+
+    _, ax = qml.draw_mpl(my_circuit)()
+
+    assert len(ax.lines) == 1  # single wire used in tape
+    assert len(ax.patches) == 1  # single pauli x gate
+    assert len(ax.texts) == 2  # one wire label, one gate label
+
+    plt.close()
diff --git a/tests/drawer/test_tape_mpl.py b/tests/drawer/test_tape_mpl.py
index 6afcc70b7..5c3699c1c 100644
--- a/tests/drawer/test_tape_mpl.py
+++ b/tests/drawer/test_tape_mpl.py
@@ -372,6 +372,8 @@ class TestSpecialGates:
         for i in range(3):
             assert ax.patches[i].center == (layer, i)
 
+        plt.close()
+
     def test_Barrier(self):
         """Test Barrier gets correct special call."""
 
@@ -543,6 +545,8 @@ class TestControlledGates:
         tape = QuantumScript.from_queue(q_tape)
         self.check_tape_controlled_qubit_unitary(tape)
 
+        plt.close()
+
     def test_nested_control_values_bool(self):
         """Test control_values get displayed correctly for nested controlled operations
         when they are provided as a list of bools."""
@@ -651,6 +655,8 @@ class TestGeneralOperations:
         assert ax.patches[0].get_width() == self.width
         assert ax.patches[0].get_height() == 2 + self.width
 
+        plt.close()
+
     @pytest.mark.parametrize("op", general_op_data)
     def test_general_operations_decimals(self, op):
         """Check that the decimals argument affects text strings when applicable."""
@@ -858,6 +864,7 @@ class TestClassicalControl:
             "linewidth": 3 * plt.rcParams["lines.linewidth"],
             "foreground": "white",  # figure.facecolor for black white sytle
         }
+        plt.close()
 
     def test_combo_measurement(self):
         """Test a control that depends on two mid circuit measurements."""
@@ -897,6 +904,8 @@ class TestClassicalControl:
         assert eraser.get_color() == plt.rcParams["figure.facecolor"]
         assert eraser.get_linewidth() == 3 * plt.rcParams["lines.linewidth"]
 
+        plt.close()
+
     def test_combo_measurement_non_terminal(self):
         """Test a combination measurement where the classical wires continue on.
         This covers the "erase_right=True" case.
@@ -937,6 +946,8 @@ class TestClassicalControl:
         assert eraser.get_color() == plt.rcParams["figure.facecolor"]
         assert eraser.get_linewidth() == 3 * plt.rcParams["lines.linewidth"]
 
+        plt.close()
+
     def test_single_mcm_measure(self):
         """Test a final measurement of a mid circuit measurement."""
 
@@ -1002,3 +1013,5 @@ class TestClassicalControl:
         assert (
             final_measure_box.get_height() == 0.75 - 2 * 0.2 + 2 * 0.25
         )  # box_length - 2 * pad + 2 *cwire_scaling
+
+        plt.close()