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:

Resolving patch:
diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -279,7 +279,20 @@ def execute(self, circuit, **kwargs):
         # generate computational basis samples
         sample_type = (SampleMP, CountsMP, ClassicalShadowMP, ShadowExpvalMP)
         if self.shots is not None or any(isinstance(m, sample_type) for m in circuit.measurements):
+            # Lightning does not support apply(rotations) anymore, so we need to rotate here
+            # Lightning without binaries fallbacks to `QubitDevice`, and hence the _CPP_BINARY_AVAILABLE condition
+            is_lightning = (
+                hasattr(self, "name")
+                and isinstance(self.name, str)
+                and "Lightning" in self.name
+                and getattr(self, "_CPP_BINARY_AVAILABLE", False)
+            )
+            diagonalizing_gates = self._get_diagonalizing_gates(circuit) if is_lightning else None
+            if is_lightning and diagonalizing_gates:
+                self.apply(diagonalizing_gates)
             self._samples = self.generate_samples()
+            if is_lightning and diagonalizing_gates:
+                self.apply([qml.adjoint(g, lazy=False) for g in reversed(diagonalizing_gates)])
 
         # compute the required statistics
         if self._shot_vector is not None:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -15,7 +15,7 @@
 """Setup file for package installation."""
 # pylint: disable=unspecified-encoding, consider-using-with
 
-from setuptools import setup, find_packages
+from setuptools import find_packages, setup
 
 with open("pennylane/_version.py") as f:
     version = f.readlines()[-1].split()[-1].strip("\"'")


Test patch:
diff --git a/tests/devices/test_lightning_qubit.py b/tests/devices/test_lightning_qubit.py
--- a/tests/devices/test_lightning_qubit.py
+++ b/tests/devices/test_lightning_qubit.py
@@ -15,6 +15,8 @@
 Tests for the accessibility of the Lightning-Qubit device
 """
 import pytest
+from flaky import flaky
+
 import pennylane as qml
 from pennylane import numpy as np
 
@@ -77,6 +79,24 @@ def circuit():
         qml.QNode(circuit, dev, diff_method="adjoint")
 
 
+@flaky(max_runs=5)
+def test_finite_shots():
+    """Test that shots in LQ and DQ give the same results."""
+
+    dev = qml.device("lightning.qubit", wires=2, shots=10000)
+    dq = qml.device("default.qubit", shots=10000)
+
+    def circuit():
+        qml.RX(np.pi / 4, 0)
+        qml.RY(-np.pi / 4, 1)
+        return qml.expval(qml.PauliY(0))
+
+    circ0 = qml.QNode(circuit, dev, diff_method=None)
+    circ1 = qml.QNode(circuit, dq, diff_method=None)
+
+    assert np.allclose(circ0(), circ1(), rtol=0.01)
+
+
 class TestDtypePreserved:
     """Test that the user-defined dtype of the device is preserved for QNode
     evaluation"""