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/tape/qscript.py b/pennylane/tape/qscript.py
--- a/pennylane/tape/qscript.py
+++ b/pennylane/tape/qscript.py
@@ -759,19 +759,8 @@ def shape(self, device):
             >>> qs.shape(dev)
             ((4,), (), (4,))
         """
-
-        if isinstance(device, qml.devices.Device):
-            # MP.shape (called below) takes 2 arguments: `device` and `shots`.
-            # With the new device interface, shots are stored on tapes rather than the device
-            # TODO: refactor MP.shape to accept `wires` instead of device (not currently done
-            # because probs.shape uses device.cutoff)
-            shots = self.shots
-        else:
-            shots = (
-                Shots(device._raw_shot_sequence)
-                if device.shot_vector is not None
-                else Shots(device.shots)
-            )
+        shots = self.shots
+        # even with the legacy device interface, the shots on the tape will agree with the shots used by the device for the execution
 
         if len(shots.shot_vector) > 1 and self.batch_size is not None:
             raise NotImplementedError(
diff --git a/pennylane/workflow/interfaces/jax_jit.py b/pennylane/workflow/interfaces/jax_jit.py
--- a/pennylane/workflow/interfaces/jax_jit.py
+++ b/pennylane/workflow/interfaces/jax_jit.py
@@ -146,20 +146,19 @@ def _execute_wrapper_inner(params, tapes, execute_fn, _, device, is_vjp=False) -
 
     def pure_callback_wrapper(p):
         new_tapes = _set_fn(tapes.vals, p)
-        res = tuple(_to_jax(execute_fn(new_tapes)))
-        # When executed under `jax.vmap` the `result_shapes_dtypes` will contain
-        # the shape without the vmap dimensions, while the function here will be
-        # executed with objects containing the vmap dimensions. So res[i].ndim
-        # will have an extra dimension for every `jax.vmap` wrapping this execution.
-        #
-        # The execute_fn will return an object with shape `(n_observables, batches)`
-        # but the default behaviour for `jax.pure_callback` is to add the extra
-        # dimension at the beginning, so `(batches, n_observables)`. So in here
-        # we detect with the heuristic above if we are executing under vmap and we
-        # swap the order in that case.
-        return jax.tree_map(lambda r, s: r.T if r.ndim > s.ndim else r, res, shape_dtype_structs)
-
-    out = jax.pure_callback(pure_callback_wrapper, shape_dtype_structs, params, vectorized=True)
+        return _to_jax(execute_fn(new_tapes))
+
+    if isinstance(device, qml.Device):
+        device_supports_vectorization = device.capabilities().get("supports_broadcasting")
+    else:
+        # first order way of determining native parameter broadcasting support
+        # will be inaccurate when inclusion of broadcast_expand depends on ExecutionConfig values (like adjoint)
+        device_supports_vectorization = (
+            qml.transforms.broadcast_expand not in device.preprocess()[0]
+        )
+    out = jax.pure_callback(
+        pure_callback_wrapper, shape_dtype_structs, params, vectorized=device_supports_vectorization
+    )
     return out
 
 


Test patch:
diff --git a/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py
--- a/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py
@@ -35,6 +35,9 @@
     [DefaultQubit(), "spsa", False, False],
     [DefaultQubit(), "hadamard", False, False],
     # [qml.device("lightning.qubit", wires=5), "adjoint", False, True],
+    [qml.device("lightning.qubit", wires=5), "adjoint", True, False],
+    [qml.device("lightning.qubit", wires=5), "adjoint", False, False],
+    [qml.device("lightning.qubit", wires=5), "parameter-shift", False, False],
 ]
 interface_and_qubit_device_and_diff_method = [
     ["auto"] + inner_list for inner_list in qubit_device_and_diff_method
@@ -78,7 +81,7 @@ def circuit(a):
             qml.RX(0.2, wires=0)
             return qml.expval(qml.PauliZ(0))
 
-        a = np.array(0.1, requires_grad=True)
+        a = jax.numpy.array(0.1)
         jax.jit(circuit)(a)
 
         assert circuit.interface == interface
@@ -299,7 +302,7 @@ def test_diff_expval_expval(
         """Test jacobian calculation"""
         if dev.name == "param_shift.qubit":
             pytest.xfail("gradient transforms have a different vjp shape convention.")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning does not support device vjps with jax jacobians.")
         gradient_kwargs = {}
 
@@ -364,7 +367,7 @@ def test_jacobian_no_evaluate(
         """Test jacobian calculation when no prior circuit evaluation has been performed"""
         if dev.name == "param_shift.qubit":
             pytest.xfail("gradient transforms have a different vjp shape convention.")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning does not support device vjps with jax jacobians.")
         gradient_kwargs = {}
 
@@ -429,7 +432,7 @@ def test_diff_single_probs(
         with a single prob output"""
         if dev.name == "param_shift.qubit":
             pytest.xfail("gradient transforms have a different vjp shape convention.")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning does not support device vjps with jax jacobians.")
         gradient_kwargs = {}
         if diff_method == "spsa":
@@ -482,7 +485,7 @@ def test_diff_multi_probs(
         with multiple prob outputs"""
         if dev.name == "param_shift.qubit":
             pytest.xfail("gradient transforms have a different vjp shape convention.")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning does not support device vjps with jax jacobians.")
         gradient_kwargs = {}
         if diff_method == "spsa":
@@ -568,7 +571,7 @@ def test_diff_expval_probs(
         with prob and expval outputs"""
         if dev.name == "param_shift.qubit":
             pytest.xfail("gradient transforms have a different vjp shape convention.")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning does not support device vjps with jax jacobians.")
         gradient_kwargs = {}
         if diff_method == "spsa":
@@ -644,7 +647,7 @@ def test_diff_expval_probs_sub_argnums(
         trainable parameters (argnums) than parameters."""
         if dev.name == "param_shift.qubit":
             pytest.xfail("gradient transforms have a different vjp shape convention.")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning does not support device vjps with jax jacobians.")
         kwargs = {}
         if diff_method == "spsa":
@@ -697,7 +700,7 @@ def test_diff_var_probs(self, dev, diff_method, grad_on_execution, device_vjp, i
         with prob and variance outputs"""
         if dev.name == "param_shift.qubit":
             pytest.xfail("gradient transforms have a different vjp shape convention.")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning does not support device vjps with jax jacobians.")
         gradient_kwargs = {}
         if diff_method == "hadamard":
@@ -1024,6 +1027,8 @@ def test_postselection_differentiation(
             pytest.skip("Diff method does not support postselection.")
         if dev.name == "param_shift.qubit":
             pytest.xfail("gradient transforms have a different vjp shape convention")
+        elif dev.name == "lightning.qubit":
+            pytest.xfail("lightning qubit does not support postselection.")
 
         @qml.qnode(
             dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution
@@ -1363,6 +1368,9 @@ def circuit(a, b):
     def test_state(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol):
         """Test that the state can be returned and differentiated"""
 
+        if dev.name == "lightning.qubit" and diff_method == "adjoint":
+            pytest.xfail("lightning.qubit does not support adjoint with the state.")
+
         x = jax.numpy.array(0.543)
         y = jax.numpy.array(-0.654)
         if not dev.wires:
@@ -1562,7 +1570,7 @@ def circuit(data, weights, coeffs):
 
     @pytest.mark.parametrize("max_diff", [1, 2])
     def test_hamiltonian_finite_shots(
-        self, dev, diff_method, grad_on_execution, device_vjp, interface, max_diff, mocker
+        self, dev, diff_method, grad_on_execution, device_vjp, interface, max_diff
     ):
         """Test that the Hamiltonian is correctly measured if there
         are non-commuting groups and the number of shots is finite
@@ -1579,7 +1587,6 @@ def test_hamiltonian_finite_shots(
             gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA, "h": H_FOR_SPSA, "num_directions": 20}
             tol = TOL_FOR_SPSA
 
-        spy = mocker.spy(qml.transforms, "hamiltonian_expand")
         obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)]
 
         @qnode(
@@ -1607,7 +1614,6 @@ def circuit(data, weights, coeffs):
         res = circuit(d, w, c, shots=50000)  # pylint:disable=unexpected-keyword-arg
         expected = c[2] * np.cos(d[1] + w[1]) - c[1] * np.sin(d[0] + w[0]) * np.sin(d[1] + w[1])
         assert np.allclose(res, expected, atol=tol)
-        spy.assert_not_called()
 
         # test gradients
         grad = jax.grad(circuit, argnums=[1, 2])(d, w, c, shots=50000)
@@ -1639,16 +1645,14 @@ def test_vmap_compared_param_broadcasting(
     ):
         """Test that jax.vmap works just as well as parameter-broadcasting with JAX JIT on the forward pass when
         vectorized=True is specified for the callback when caching is disabled."""
+        if (
+            dev.name == "default.qubit"
+            and diff_method == "adjoint"
+            and grad_on_execution
+            and not device_vjp
+        ):
+            pytest.xfail("adjoint is incompatible with parameter broadcasting.")
         interface = "jax-jit"
-        if diff_method == "adjoint":
-            pytest.skip("The adjoint method does not yet support Hamiltonians")
-        elif diff_method == "hadamard":
-            pytest.skip("The Hadamard method does not yet support Hamiltonians")
-
-        if diff_method == "backprop":
-            pytest.skip(
-                "The backprop method does not yet support parameter-broadcasting with Hamiltonians"
-            )
 
         n_configs = 5
         pars_q = np.random.rand(n_configs, 2)
@@ -1665,15 +1669,14 @@ def minimal_circ(params):
             def _measure_operator():
                 qml.RY(params[..., 0], wires=0)
                 qml.RY(params[..., 1], wires=1)
-                op = qml.Hamiltonian([1.0], [qml.PauliZ(0) @ qml.PauliZ(1)])
-                return qml.expval(op)
+                return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
 
             res = _measure_operator()
             return res
 
-        assert np.allclose(
-            jax.jit(minimal_circ)(pars_q), jax.jit(jax.vmap(minimal_circ))(pars_q), tol
-        )
+        res1 = jax.jit(minimal_circ)(pars_q)
+        res2 = jax.jit(jax.vmap(minimal_circ))(pars_q)
+        assert np.allclose(res1, res2, tol)
 
     def test_vmap_compared_param_broadcasting_multi_output(
         self, dev, diff_method, grad_on_execution, device_vjp, interface, tol
@@ -1681,16 +1684,56 @@ def test_vmap_compared_param_broadcasting_multi_output(
         """Test that jax.vmap works just as well as parameter-broadcasting with JAX JIT on the forward pass when
         vectorized=True is specified for the callback when caching is disabled and when multiple output values
         are returned."""
+        if (
+            dev.name == "default.qubit"
+            and diff_method == "adjoint"
+            and grad_on_execution
+            and not device_vjp
+        ):
+            pytest.xfail("adjoint is incompatible with parameter broadcasting.")
         interface = "jax-jit"
-        if diff_method == "adjoint":
-            pytest.skip("The adjoint method does not yet support Hamiltonians")
-        elif diff_method == "hadamard":
-            pytest.skip("The Hadamard method does not yet support Hamiltonians")
 
-        if diff_method == "backprop":
-            pytest.skip(
-                "The backprop method does not yet support parameter-broadcasting with Hamiltonians"
+        n_configs = 5
+        pars_q = np.random.rand(n_configs, 2)
+
+        def minimal_circ(params):
+            @qml.qnode(
+                dev,
+                interface=interface,
+                diff_method=diff_method,
+                grad_on_execution=grad_on_execution,
+                device_vjp=device_vjp,
+                cache=None,
             )
+            def _measure_operator():
+                qml.RY(params[..., 0], wires=0)
+                qml.RY(params[..., 1], wires=1)
+                return qml.expval(qml.Z(0) @ qml.Z(1)), qml.expval(qml.X(0) @ qml.X(1))
+
+            res = _measure_operator()
+            return res
+
+        res1, res2 = jax.jit(minimal_circ)(pars_q)
+        vres1, vres2 = jax.jit(jax.vmap(minimal_circ))(pars_q)
+        assert np.allclose(res1, vres1, tol)
+        assert np.allclose(res2, vres2, tol)
+
+    def test_vmap_compared_param_broadcasting_probs(
+        self, dev, diff_method, grad_on_execution, device_vjp, interface, tol
+    ):
+        """Test that jax.vmap works just as well as parameter-broadcasting with JAX JIT on the forward pass when
+        vectorized=True is specified for the callback when caching is disabled and when multiple output values
+        are returned."""
+        if (
+            dev.name == "default.qubit"
+            and diff_method == "adjoint"
+            and grad_on_execution
+            and not device_vjp
+        ):
+            pytest.xfail("adjoint is incompatible with parameter broadcasting.")
+        elif dev.name == "lightning.qubit" and diff_method == "adjoint":
+            pytest.xfail("lightning adjoign cannot differentiate probabilities.")
+        interface = "jax-jit"
 
         n_configs = 5
         pars_q = np.random.rand(n_configs, 2)
@@ -1707,9 +1750,7 @@ def minimal_circ(params):
             def _measure_operator():
                 qml.RY(params[..., 0], wires=0)
                 qml.RY(params[..., 1], wires=1)
-                op1 = qml.Hamiltonian([1.0], [qml.PauliZ(0) @ qml.PauliZ(1)])
-                op2 = qml.Hamiltonian([1.0], [qml.PauliX(0) @ qml.PauliX(1)])
-                return qml.expval(op1), qml.expval(op2)
+                return qml.probs(wires=0), qml.probs(wires=1)
 
             res = _measure_operator()
             return res
@@ -1739,7 +1780,7 @@ def test_gradient(
         gradient_kwargs = {}
         if dev.name == "param_shift.qubit":
             pytest.xfail("gradient transforms have a different vjp shape convention.")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device vjps are not compatible with jax jaocbians")
         if device_vjp and jacobian == jax.jacfwd:
             pytest.skip("device vjps not compatible with forward diff.")
@@ -1842,7 +1883,7 @@ def test_gradient_subset(
         to a subset of arguments"""
         if diff_method == "spsa" and not grad_on_execution and not device_vjp:
             pytest.xfail(reason="incorrect jacobian results")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device vjps are not compatible with jax jaocbians")
 
         if diff_method == "device" and not grad_on_execution and device_vjp:
@@ -1883,7 +1924,7 @@ def test_gradient_scalar_cost_vector_valued_qnode(
         gradient_kwargs = {}
         if dev.name == "param_shift.qubit":
             pytest.xfail("gradient transforms have a different vjp shape convention.")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device vjps are not compatible with jax jaocbians")
         elif jacobian == jax.jacfwd and device_vjp:
             pytest.skip("device vjps are not compatible with forward differentiation.")
@@ -1937,7 +1978,7 @@ def test_matrix_parameter(
         parameter"""
         if dev.name == "param_shift.qubit":
             pytest.xfail("gradient transforms have a different vjp shape convention.")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device vjps are not compatible with jax jaocbians")
         if jacobian == jax.jacfwd and device_vjp:
             pytest.skip("device vjps are not compatible with forward differentiation.")
@@ -1980,7 +2021,7 @@ def test_grad_single_measurement_param(
         """For one measurement and one param, the gradient is a float."""
         if dev.name == "param_shift.qubit":
             pytest.xfail("gradient transforms have a different vjp shape convention")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device vjps are not compatible with jax jaocbians")
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
@@ -2016,7 +2057,7 @@ def test_grad_single_measurement_multiple_param(
             pytest.skip("Test does not support finite shots and adjoint/backprop")
         if jacobian == jax.jacfwd and device_vjp:
             pytest.skip("jacfwd is not compatible with device_vjp=True.")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device vjps are not compatible with jax jaocbians")
 
         @qnode(
@@ -2049,7 +2090,7 @@ def test_grad_single_measurement_multiple_param_array(
         """For one measurement and multiple param as a single array params, the gradient is an array."""
         if dev.name == "param_shift.qubit":
             pytest.xfail("gradient transforms have a different vjp shape convention.")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device vjps are not compatible with jax jaocbians")
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
@@ -2082,7 +2123,7 @@ def test_jacobian_single_measurement_param_probs(
         dimension"""
         if dev.name == "param_shift.qubit":
             pytest.xfail("gradient transforms have a different vjp shape convention.")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device vjps are not compatible with jax jaocbians")
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
@@ -2115,7 +2156,7 @@ def test_jacobian_single_measurement_probs_multiple_param(
         the correct dimension"""
         if dev.name == "param_shift.qubit":
             pytest.xfail("gradient transforms have a different vjp shape convention.")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device vjps are not compatible with jax jaocbians")
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
@@ -2156,7 +2197,7 @@ def test_jacobian_single_measurement_probs_multiple_param_single_array(
             pytest.xfail("gradient transforms have a different vjp shape convention.")
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device vjps are not compatible with jax jaocbians")
         if jacobian == jax.jacfwd and device_vjp:
             pytest.skip("jacfwd is not compatible with device_vjp=True.")
@@ -2185,7 +2226,7 @@ def test_jacobian_expval_expval_multiple_params(
         """The jacobian of multiple measurements with multiple params return a tuple of arrays."""
         if dev.name == "param_shift.qubit":
             pytest.xfail("gradient transforms have a different vjp shape convention.")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device vjps are not compatible with jax jaocbians")
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
@@ -2234,7 +2275,7 @@ def test_jacobian_expval_expval_multiple_params_array(
         """The jacobian of multiple measurements with a multiple params array return a single array."""
         if dev.name == "param_shift.qubit":
             pytest.xfail("gradient transforms have a different vjp shape convention.")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device vjps are not compatible with jax jaocbians")
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
@@ -2364,7 +2405,7 @@ def test_jacobian_multiple_measurement_single_param(
         """The jacobian of multiple measurements with a single params return an array."""
         if dev.name == "param_shift.qubit":
             pytest.xfail("gradient transforms have a different vjp shape convention")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device vjps are not compatible with jax jaocbians")
         if device_vjp and jacobian == jax.jacfwd:
             pytest.skip("device vjp not compatible with forward differentiation.")
@@ -2402,7 +2443,7 @@ def test_jacobian_multiple_measurement_multiple_param(
         """The jacobian of multiple measurements with a multiple params return a tuple of arrays."""
         if dev.name == "param_shift.qubit":
             pytest.xfail("gradient transforms have a different vjp shape convention.")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device vjps are not compatible with jax jaocbians")
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
@@ -2449,7 +2490,7 @@ def test_jacobian_multiple_measurement_multiple_param_array(
         """The jacobian of multiple measurements with a multiple params array return a single array."""
         if dev.name == "param_shift.qubit":
             pytest.xfail("gradient transforms have a different vjp shape convention.")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device vjps are not compatible with jax jaocbians")
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
@@ -2853,7 +2894,7 @@ def test_single_measurement(
         kwargs = {}
         if jacobian == jax.jacfwd and device_vjp:
             pytest.skip("jacfwd is not compatible with device_vjp=True.")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device vjps are not compatible with jax jaocbians")
         if dev.name == "param_shift.qubit":
             pytest.xfail("gradient transform have a different vjp shape convention.")
@@ -2909,7 +2950,7 @@ def test_multi_measurements(
         """Test multiple measurements with different diff methods with argnums."""
         if jacobian == jax.jacfwd and device_vjp:
             pytest.skip("jacfwd is not compatible with device_vjp=True.")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device vjps are not compatible with jax jaocbians")
         if dev.name == "param_shift.qubit":
             pytest.xfail("gradient transform have a different vjp shape convention.")
diff --git a/tests/interfaces/default_qubit_2_integration/test_jax_qnode_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_jax_qnode_default_qubit_2.py
--- a/tests/interfaces/default_qubit_2_integration/test_jax_qnode_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_jax_qnode_default_qubit_2.py
@@ -277,7 +277,7 @@ def test_diff_expval_expval(
         if diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
             tol = TOL_FOR_SPSA
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device_vjp not compatible with jax.jacobian.")
 
         a = jax.numpy.array(0.1)
@@ -333,7 +333,7 @@ def test_jacobian_no_evaluate(
             "device_vjp": device_vjp,
         }
 
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device_vjp not compatible with jax.jacobian.")
 
         if diff_method == "spsa":
@@ -392,7 +392,7 @@ def test_diff_single_probs(
         if diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
             tol = TOL_FOR_SPSA
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device_vjp not compatible with jax.jacobian.")
 
         x = jax.numpy.array(0.543)
@@ -441,7 +441,7 @@ def test_diff_multi_probs(
         if diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
             tol = TOL_FOR_SPSA
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device_vjp not compatible with jax.jacobian.")
 
         x = jax.numpy.array(0.543)
@@ -522,7 +522,7 @@ def test_diff_expval_probs(
         if diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
             tol = TOL_FOR_SPSA
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device_vjp not compatible with jax.jacobian.")
 
         x = jax.numpy.array(0.543)
@@ -609,7 +609,7 @@ def circuit(x, y):
             qml.CNOT(wires=[0, 1])
             return qml.expval(qml.PauliZ(0)), qml.probs(wires=[1])
 
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning does not support measuring probabilities with adjoint.")
         jac = jax.jacobian(circuit, argnums=[0])(x, y)
 
@@ -647,7 +647,7 @@ def test_diff_var_probs(self, dev, diff_method, grad_on_execution, device_vjp, i
 
         if diff_method == "hadamard":
             pytest.skip("Hadamard does not support var")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device_vjp not compatible with jax.jacobian.")
         elif diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
@@ -1712,7 +1712,7 @@ def test_jacobian_expval_expval_multiple_params(
             pytest.skip("Test does not support finite shots and adjoint/backprop")
         if device_vjp and jacobian is jax.jacfwd:
             pytest.skip("forward pass can't be done with registered vjp.")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device_vjp not compatible with jax.jacobian.")
         par_0 = jax.numpy.array(0.1)
         par_1 = jax.numpy.array(0.2)
@@ -1758,7 +1758,7 @@ def test_jacobian_expval_expval_multiple_params_array(
             pytest.skip("Test does not support finite shots and adjoint/backprop")
         if device_vjp and jacobian is jax.jacfwd:
             pytest.skip("forward pass can't be done with registered vjp.")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device_vjp not compatible with jax.jacobian.")
 
         @qnode(
@@ -1878,7 +1878,7 @@ def test_jacobian_multiple_measurement_single_param(
         """The jacobian of multiple measurements with a single params return an array."""
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device_vjp not compatible with jax.jacobian.")
         if diff_method == "adjoint" and jacobian == jax.jacfwd:
             pytest.skip("jacfwd doesn't like complex numbers")
@@ -1915,7 +1915,7 @@ def test_jacobian_multiple_measurement_multiple_param(
         """The jacobian of multiple measurements with a multiple params return a tuple of arrays."""
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device_vjp not compatible with jax.jacobian.")
         if diff_method == "adjoint" and jacobian == jax.jacfwd:
             pytest.skip("jacfwd doesn't like complex numbers")
@@ -1961,7 +1961,7 @@ def test_jacobian_multiple_measurement_multiple_param_array(
         """The jacobian of multiple measurements with a multiple params array return a single array."""
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
-        if "lightning" in getattr(dev, "short_name", ""):
+        if "lightning" in dev.name:
             pytest.xfail("lightning device_vjp not compatible with jax.jacobian.")
         if diff_method == "adjoint" and jacobian == jax.jacfwd:
             pytest.skip("jacfwd doesn't like complex numbers")
diff --git a/tests/interfaces/test_jax_jit.py b/tests/interfaces/test_jax_jit.py
--- a/tests/interfaces/test_jax_jit.py
+++ b/tests/interfaces/test_jax_jit.py
@@ -811,7 +811,7 @@ def cost(a, cache):
                 qml.RY(a[2], wires=0)
                 qml.sample(qml.PauliZ(0))
 
-            tape = qml.tape.QuantumScript.from_queue(q)
+            tape = qml.tape.QuantumScript.from_queue(q, shots=dev.shots)
 
             res = qml.execute([tape], dev, cache=cache, **execute_kwargs)[0]
             return res