Repository: PennyLaneAI/pennylane

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.


Discussion:
Thanks for opening this issue.  It looks like we were only applying the transform program when the device inherited from the `qml.devices.Device`.

https://github.com/PennyLaneAI/pennylane/blob/d7593a30739910dab72f09e05130cda029f57de8/pennylane/drawer/draw.py#L269

On the plus side, the correctly transformed tape was getting to the device. The transformed version just wasn't showing up in the drawer.  I should be able to put together a fix rather quickly.

Also note that we will be improving the entire system for where in the workflow we draw circuits in #5139 . This will allow much more control over which transforms are applied before drawing.

Resolving patch:
diff --git a/pennylane/drawer/draw.py b/pennylane/drawer/draw.py
--- a/pennylane/drawer/draw.py
+++ b/pennylane/drawer/draw.py
@@ -259,16 +259,16 @@ def wrapper(*args, **kwargs):
             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 wrapper(*args, **kwargs):
                     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 wrapper(*args, **kwargs_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
--- a/tests/drawer/test_draw.py
+++ b/tests/drawer/test_draw.py
@@ -881,6 +881,29 @@ def circ(t):
     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
--- a/tests/drawer/test_draw_mpl.py
+++ b/tests/drawer/test_draw_mpl.py
@@ -398,3 +398,31 @@ def circuit():
     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
--- a/tests/drawer/test_tape_mpl.py
+++ b/tests/drawer/test_tape_mpl.py
@@ -372,6 +372,8 @@ def test_CCZ(self):
         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 @@ def test_control_values_bool(self):
         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 @@ def test_snapshot(self):
         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 @@ def test_single_measure_multiple_conds(self):
             "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 @@ def test_combo_measurement(self):
         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 @@ def test_combo_measurement_non_terminal(self):
         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 @@ def test_multiple_mcm_measure(self):
         assert (
             final_measure_box.get_height() == 0.75 - 2 * 0.2 + 2 * 0.25
         )  # box_length - 2 * pad + 2 *cwire_scaling
+
+        plt.close()