Repository: PennyLaneAI/pennylane

Issue:
[BUG] JIT + Global measurements without device wires does not work
### Expected behavior

Inferring the wires of a global measurement from the tape's operations if the device has no wires specified works without and with JITting.

### Actual behavior

Only works without JITting.

### Additional information

@albi3ro suggested to modify `MeasurementProcess.shape` to take `num_device_wires` as an argument, instead of the full device.

### Source code

```shell
dev = qml.device("default.qubit")

@jax.jit
@qml.qnode(dev, diff_method="parameter-shift")
def node(x):
    qml.RX(x,0)
    return qml.probs() # Crashes, see thread
    return qml.probs(wires=[0, 1]) # Works

node(jnp.array(0.5))
```


### Tracebacks

```shell
---------------------------------------------------------------------------
XlaRuntimeError                           Traceback (most recent call last)
Cell In[25], line 9
      6     qml.RX(x,0)
      7     return qml.probs()
----> 9 node(jnp.array(0.5))

    [... skipping hidden 10 frame]

File ~/venvs/dev/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py:1151, in ExecuteReplicated.__call__(self, *args)
   1148 if (self.ordered_effects or self.has_unordered_effects
   1149     or self.has_host_callbacks):
   1150   input_bufs = self._add_tokens_to_inputs(input_bufs)
-> 1151   results = self.xla_executable.execute_sharded(
   1152       input_bufs, with_tokens=True
   1153   )
   1154   result_token_bufs = results.disassemble_prefix_into_single_device_arrays(
   1155       len(self.ordered_effects))
   1156   sharded_runtime_token = results.consume_token()

XlaRuntimeError: INTERNAL: Generated function failed: CpuCallback error: RuntimeError: Incorrect output shape for return value 0: Expected: (1,), Actual: (2,)

At:
  /home/david/venvs/dev/lib/python3.10/site-packages/jax/_src/interpreters/mlir.py(2374): _wrapped_callback
  /home/david/venvs/dev/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py(1151): __call__
  /home/david/venvs/dev/lib/python3.10/site-packages/jax/_src/profiler.py(336): wrapper
  /home/david/venvs/dev/lib/python3.10/site-packages/jax/_src/pjit.py(1185): _pjit_call_impl_python
  /home/david/venvs/dev/lib/python3.10/site-packages/jax/_src/pjit.py(1229): call_impl_cache_miss
  /home/david/venvs/dev/lib/python3.10/site-packages/jax/_src/pjit.py(1245): _pjit_call_impl
  /home/david/venvs/dev/lib/python3.10/site-packages/jax/_src/core.py(935): process_primitive
  /home/david/venvs/dev/lib/python3.10/site-packages/jax/_src/core.py(447): bind_with_trace
  /home/david/venvs/dev/lib/python3.10/site-packages/jax/_src/core.py(2740): bind
  /home/david/venvs/dev/lib/python3.10/site-packages/jax/_src/pjit.py(168): _python_pjit_helper
  /home/david/venvs/dev/lib/python3.10/site-packages/jax/_src/pjit.py(257): cache_miss
  /home/david/venvs/dev/lib/python3.10/site-packages/jax/_src/traceback_util.py(179): reraise_with_filtered_traceback
  /tmp/ipykernel_292987/2553214436.py(9): <module>
  /home/david/venvs/dev/lib/python3.10/site-packages/IPython/core/interactiveshell.py(3577): run_code
  /home/david/venvs/dev/lib/python3.10/site-packages/IPython/core/interactiveshell.py(3517): run_ast_nodes
  /home/david/venvs/dev/lib/python3.10/site-packages/IPython/core/interactiveshell.py(3334): run_cell_async
  /home/david/venvs/dev/lib/python3.10/site-packages/IPython/core/async_helpers.py(129): _pseudo_sync_runner
  /home/david/venvs/dev/lib/python3.10/site-packages/IPython/core/interactiveshell.py(3130): _run_cell
  /home/david/venvs/dev/lib/python3.10/site-packages/IPython/core/interactiveshell.py(3075): run_cell
  /home/david/venvs/dev/lib/python3.10/site-packages/ipykernel/zmqshell.py(549): run_cell
  /home/david/venvs/dev/lib/python3.10/site-packages/ipykernel/ipkernel.py(449): do_execute
  /home/david/venvs/dev/lib/python3.10/site-packages/ipykernel/kernelbase.py(778): execute_request
  /home/david/venvs/dev/lib/python3.10/site-packages/ipykernel/ipkernel.py(362): execute_request
  /home/david/venvs/dev/lib/python3.10/site-packages/ipykernel/kernelbase.py(437): dispatch_shell
  /home/david/venvs/dev/lib/python3.10/site-packages/ipykernel/kernelbase.py(534): process_one
  /home/david/venvs/dev/lib/python3.10/site-packages/ipykernel/kernelbase.py(545): dispatch_queue
  /usr/lib/python3.10/asyncio/events.py(80): _run
  /usr/lib/python3.10/asyncio/base_events.py(1909): _run_once
  /usr/lib/python3.10/asyncio/base_events.py(603): run_forever
  /home/david/venvs/dev/lib/python3.10/site-packages/tornado/platform/asyncio.py(205): start
  /home/david/venvs/dev/lib/python3.10/site-packages/ipykernel/kernelapp.py(739): start
  /home/david/venvs/dev/lib/python3.10/site-packages/traitlets/config/application.py(1075): launch_instance
  /home/david/venvs/dev/lib/python3.10/site-packages/ipykernel_launcher.py(18): <module>
  /usr/lib/python3.10/runpy.py(86): _run_code
  /usr/lib/python3.10/runpy.py(196): _run_module_as_main
```


### System information

```shell
pl dev
```


### Existing GitHub issues

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

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6108?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.65%. Comparing base [(`110c00c`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/110c00c122896a79e33c7cfc1210b75db2dd8050?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`187bcf0`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/187bcf05bc291eb377232eb14ce88f84e09e7ac1?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 1 commits behind head on master.

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


```diff
@@            Coverage Diff             @@
##           master    #6108      +/-   ##
==========================================
- Coverage   99.66%   99.65%   -0.01%     
==========================================
  Files         432      432              
  Lines       41988    41655     -333     
==========================================
- Hits        41847    41513     -334     
- Misses        141      142       +1     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6108?dropdown=coverage&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/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index a67925c4a..b60a5b400 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -40,6 +40,11 @@
 
 <h3>Improvements 🛠</h3>
 
+* Counts measurements with `all_outcomes=True` can now be used with jax jitting. Measurements
+  broadcasted across all available wires (`qml.probs()`) can now be used with jit and devices that
+  allow variable numbers of wires (`qml.device('default.qubit')`).
+  [(#6108)](https://github.com/PennyLaneAI/pennylane/pull/6108/)
+
 <h4>A Prep-Select-Prep template</h4>
 
 * The `qml.PrepSelPrep` template is added. The template implements a block-encoding of a linear
@@ -270,6 +275,15 @@
 
 <h3>Breaking changes 💔</h3>
 
+* `MeasurementProcess.shape(shots: Shots, device:Device)` is now
+  `MeasurementProcess.shape(shots: Optional[int], num_device_wires:int = 0)`. This has been done to allow
+  jitting when a measurement is broadcasted across all available wires, but the device does not specify wires.
+  [(#6108)](https://github.com/PennyLaneAI/pennylane/pull/6108/)
+
+* If the shape of a probability measurement is affected by a `Device.cutoff` property, it will no longer work with
+  jitting.
+  [(#6108)](https://github.com/PennyLaneAI/pennylane/pull/6108/)
+
 * `GlobalPhase` is considered non-differentiable with tape transforms.
   As a consequence, `qml.gradients.finite_diff` and `qml.gradients.spsa_grad` no longer
   support differentiation of `GlobalPhase` with state-based outputs.
diff --git a/pennylane/devices/null_qubit.py b/pennylane/devices/null_qubit.py
index 9e52ac72e..47d695ee5 100644
--- a/pennylane/devices/null_qubit.py
+++ b/pennylane/devices/null_qubit.py
@@ -22,7 +22,7 @@ import logging
 from dataclasses import replace
 from functools import singledispatch
 from numbers import Number
-from typing import Union
+from typing import Optional, Union
 
 import numpy as np
 
@@ -37,7 +37,6 @@ from pennylane.measurements import (
     MeasurementProcess,
     MeasurementValue,
     ProbabilityMP,
-    Shots,
     StateMP,
 )
 from pennylane.tape import QuantumScriptOrBatch
@@ -54,57 +53,52 @@ logger.addHandler(logging.NullHandler())
 
 @singledispatch
 def zero_measurement(
-    mp: MeasurementProcess, obj_with_wires, shots: Shots, batch_size: int, interface: str
+    mp: MeasurementProcess, num_device_wires, shots: Optional[int], batch_size: int, interface: str
 ):
     """Create all-zero results for various measurement processes."""
-    return _zero_measurement(mp, obj_with_wires, shots, batch_size, interface)
+    return _zero_measurement(mp, num_device_wires, shots, batch_size, interface)
 
 
-def _zero_measurement(mp, obj_with_wires, shots, batch_size, interface):
-    shape = mp.shape(obj_with_wires, shots)
-    if all(isinstance(s, int) for s in shape):
-        if batch_size is not None:
-            shape = (batch_size,) + shape
-        return math.zeros(shape, like=interface, dtype=mp.numeric_type)
+def _zero_measurement(
+    mp: MeasurementProcess, num_device_wires: int, shots: Optional[int], batch_size, interface
+):
+    shape = mp.shape(shots, num_device_wires)
     if batch_size is not None:
-        shape = ((batch_size,) + s for s in shape)
-    return tuple(math.zeros(s, like=interface, dtype=mp.numeric_type) for s in shape)
+        shape = (batch_size,) + shape
+    return math.zeros(shape, like=interface, dtype=mp.numeric_type)
 
 
 @zero_measurement.register
-def _(mp: ClassicalShadowMP, obj_with_wires, shots, batch_size, interface):
-    shapes = [mp.shape(obj_with_wires, Shots(s)) for s in shots]
+def _(mp: ClassicalShadowMP, num_device_wires, shots: Optional[int], batch_size, interface):
     if batch_size is not None:
         # shapes = [(batch_size,) + shape for shape in shapes]
         raise ValueError(
             "Parameter broadcasting is not supported with null.qubit and qml.classical_shadow"
         )
-    results = tuple(math.zeros(shape, like=interface, dtype=np.int8) for shape in shapes)
-    return results if shots.has_partitioned_shots else results[0]
+    shape = mp.shape(shots, num_device_wires)
+    return math.zeros(shape, like=interface, dtype=np.int8)
 
 
 @zero_measurement.register
-def _(mp: CountsMP, obj_with_wires, shots, batch_size, interface):
+def _(mp: CountsMP, num_device_wires, shots, batch_size, interface):
     outcomes = []
     if mp.obs is None and not isinstance(mp.mv, MeasurementValue):
-        num_wires = len(obj_with_wires.wires)
-        state = "0" * num_wires
-        results = tuple({state: math.asarray(s, like=interface)} for s in shots)
+        state = "0" * num_device_wires
+        results = {state: math.asarray(shots, like=interface)}
         if mp.all_outcomes:
-            outcomes = [f"{x:0{num_wires}b}" for x in range(1, 2**num_wires)]
+            outcomes = [f"{x:0{num_device_wires}b}" for x in range(1, 2**num_device_wires)]
     else:
         outcomes = sorted(mp.eigvals())  # always assign shots to the smallest
-        results = tuple({outcomes[0]: math.asarray(s, like=interface)} for s in shots)
+        results = {outcomes[0]: math.asarray(shots, like=interface)}
         outcomes = outcomes[1:] if mp.all_outcomes else []
 
     if outcomes:
         zero = math.asarray(0, like=interface)
-        for res in results:
-            for val in outcomes:
-                res[val] = zero
+        for val in outcomes:
+            results[val] = zero
     if batch_size is not None:
-        results = tuple([r] * batch_size for r in results)
-    return results[0] if len(results) == 1 else results
+        results = tuple(results for _ in range(batch_size))
+    return results
 
 
 zero_measurement.register(DensityMatrixMP)(_zero_measurement)
@@ -112,13 +106,18 @@ zero_measurement.register(DensityMatrixMP)(_zero_measurement)
 
 @zero_measurement.register(StateMP)
 @zero_measurement.register(ProbabilityMP)
-def _(mp: Union[StateMP, ProbabilityMP], obj_with_wires, shots, batch_size, interface):
-    wires = mp.wires or obj_with_wires.wires
-    state = [1.0] + [0.0] * (2 ** len(wires) - 1)
+def _(
+    mp: Union[StateMP, ProbabilityMP],
+    num_device_wires: int,
+    shots: Optional[int],
+    batch_size,
+    interface,
+):
+    num_wires = len(mp.wires) or num_device_wires
+    state = [1.0] + [0.0] * (2**num_wires - 1)
     if batch_size is not None:
         state = [state] * batch_size
-    result = math.asarray(state, like=interface)
-    return (result,) * shots.num_copies if shots.has_partitioned_shots else result
+    return math.asarray(state, like=interface)
 
 
 @simulator_tracking
@@ -222,26 +221,26 @@ class NullQubit(Device):
         self._debugger = None
 
     def _simulate(self, circuit, interface):
-        shots = circuit.shots
-        obj_with_wires = self if self.wires else circuit
-        results = tuple(
-            zero_measurement(mp, obj_with_wires, shots, circuit.batch_size, interface)
-            for mp in circuit.measurements
-        )
-        if len(results) == 1:
-            return results[0]
-        if shots.has_partitioned_shots:
-            return tuple(zip(*results))
-        return results
+        num_device_wires = len(self.wires) if self.wires else len(circuit.wires)
+        results = []
+        for s in circuit.shots or [None]:
+            r = tuple(
+                zero_measurement(mp, num_device_wires, s, circuit.batch_size, interface)
+                for mp in circuit.measurements
+            )
+            results.append(r[0] if len(circuit.measurements) == 1 else r)
+        if circuit.shots.has_partitioned_shots:
+            return tuple(results)
+        return results[0]
 
     def _derivatives(self, circuit, interface):
         shots = circuit.shots
-        obj_with_wires = self if self.wires else circuit
+        num_device_wires = len(self.wires) if self.wires else len(circuit.wires)
         n = len(circuit.trainable_params)
         derivatives = tuple(
             (
                 math.zeros_like(
-                    zero_measurement(mp, obj_with_wires, shots, circuit.batch_size, interface)
+                    zero_measurement(mp, num_device_wires, shots, circuit.batch_size, interface)
                 ),
             )
             * n
diff --git a/pennylane/measurements/classical_shadow.py b/pennylane/measurements/classical_shadow.py
index 7ac08e7a7..6c502858c 100644
--- a/pennylane/measurements/classical_shadow.py
+++ b/pennylane/measurements/classical_shadow.py
@@ -450,9 +450,9 @@ class ClassicalShadowMP(MeasurementTransform):
     ) -> tuple:
         return (2, shots, n_wires), np.int8
 
-    def shape(self, device, shots):  # pylint: disable=unused-argument
+    def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple[int, int, int]:
         # otherwise, the return type requires a device
-        if not shots:
+        if shots is None:
             raise MeasurementShapeError(
                 "Shots must be specified to obtain the shape of a classical "
                 "shadow measurement process."
@@ -460,7 +460,7 @@ class ClassicalShadowMP(MeasurementTransform):
 
         # the first entry of the tensor represents the measured bits,
         # and the second indicate the indices of the unitaries used
-        return (2, shots.total_shots, len(self.wires))
+        return (2, shots, len(self.wires))
 
     def __copy__(self):
         return self.__class__(
@@ -559,12 +559,8 @@ class ShadowExpvalMP(MeasurementTransform):
     def return_type(self):
         return ShadowExpval
 
-    def shape(self, device, shots):
-        is_single_op = isinstance(self.H, Operator)
-        if not shots.has_partitioned_shots:
-            return () if is_single_op else (len(self.H),)
-        base = () if is_single_op else (len(self.H),)
-        return (base,) * sum(s.copies for s in shots.shot_vector)
+    def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple:
+        return () if isinstance(self.H, Operator) else (len(self.H),)
 
     @property
     def wires(self):
diff --git a/pennylane/measurements/expval.py b/pennylane/measurements/expval.py
index 95499fe9e..1a95f9e7c 100644
--- a/pennylane/measurements/expval.py
+++ b/pennylane/measurements/expval.py
@@ -98,11 +98,8 @@ class ExpectationMP(SampleMeasurement, StateMeasurement):
     def numeric_type(self):
         return float
 
-    def shape(self, device, shots):
-        if not shots.has_partitioned_shots:
-            return ()
-        num_shot_elements = sum(s.copies for s in shots.shot_vector)
-        return tuple(() for _ in range(num_shot_elements))
+    def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple:
+        return ()
 
     def process_samples(
         self,
diff --git a/pennylane/measurements/measurements.py b/pennylane/measurements/measurements.py
index b6d41dd09..ab28f4874 100644
--- a/pennylane/measurements/measurements.py
+++ b/pennylane/measurements/measurements.py
@@ -17,7 +17,6 @@ outcomes from quantum observables - expectation values, variances of expectation
 and measurement samples using AnnotatedQueues.
 """
 import copy
-import functools
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
 from enum import Enum
@@ -30,8 +29,6 @@ from pennylane.pytrees import register_pytree
 from pennylane.typing import TensorLike
 from pennylane.wires import Wires
 
-from .shots import Shots
-
 # =============================================================================
 # ObservableReturnTypes types
 # =============================================================================
@@ -289,58 +286,35 @@ class MeasurementProcess(ABC, metaclass=qml.capture.ABCCaptureMeta):
             f"The numeric type of the measurement {self.__class__.__name__} is not defined."
         )
 
-    def shape(self, device, shots: Shots) -> tuple:
-        """The expected output shape of the MeasurementProcess.
-
-        Note that the output shape is dependent on the shots or device when:
-
-        * The measurement type is either ``_Probability``, ``_State`` (from :func:`.state`) or
-          ``_Sample``;
-        * The shot vector was defined.
-
-        For example, assuming a device with ``shots=None``, expectation values
-        and variances define ``shape=(,)``, whereas probabilities in the qubit
-        model define ``shape=(2**num_wires)`` where ``num_wires`` is the
-        number of wires the measurement acts on.
+    def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple[int, ...]:
+        """Calculate the shape of the result object tensor.
 
         Args:
-            device (pennylane.Device): a PennyLane device to use for determining the shape
-            shots (~.Shots): object defining the number and batches of shots
+            shots (Optional[int]) = None: the number of shots used execute the circuit. ``None``
+               indicates an analytic simulation.  Shot vectors are handled by calling this method
+               multiple times.
+            num_device_wires (int)=0 : The number of wires that will be used if the measurement is
+               broadcasted across all available wires (``len(mp.wires) == 0``). If the device
+               itself doesn't provide a number of wires, the number of tape wires will be provided
+               here instead:
 
         Returns:
-            tuple: the output shape
+            tuple[int,...]: An arbitrary length tuple of ints.  May be an empty tuple.
+
+        >>> qml.probs(wires=(0,1)).shape()
+        (4,)
+        >>> qml.sample(wires=(0,1)).shape(shots=50)
+        (50, 2)
+        >>> qml.state().shape(num_device_wires=4)
+        (16,)
+        >>> qml.expval(qml.Z(0)).shape()
+        ()
 
-        Raises:
-            QuantumFunctionError: the return type of the measurement process is
-                unrecognized and cannot deduce the numeric type
         """
         raise qml.QuantumFunctionError(
             f"The shape of the measurement {self.__class__.__name__} is not defined"
         )
 
-    @staticmethod
-    @functools.lru_cache()
-    def _get_num_basis_states(num_wires, device):
-        """Auxiliary function to determine the number of basis states given the
-        number of systems and a quantum device.
-
-        This function is meant to be used with the Probability measurement to
-        determine how many outcomes there will be. With qubit based devices
-        we'll have two outcomes for each subsystem. With continuous variable
-        devices that impose a Fock cutoff the number of basis states per
-        subsystem equals the cutoff value.
-
-        Args:
-            num_wires (int): the number of qubits/qumodes
-            device (pennylane.Device): a PennyLane device
-
-        Returns:
-            int: the number of basis states
-        """
-        cutoff = getattr(device, "cutoff", None)
-        base = 2 if cutoff is None else cutoff
-        return base**num_wires
-
     @qml.QueuingManager.stop_recording()
     def diagonalizing_gates(self):
         """Returns the gates that diagonalize the measured wires such that they
diff --git a/pennylane/measurements/mutual_info.py b/pennylane/measurements/mutual_info.py
index 95c71f196..18335fe43 100644
--- a/pennylane/measurements/mutual_info.py
+++ b/pennylane/measurements/mutual_info.py
@@ -154,11 +154,8 @@ class MutualInfoMP(StateMeasurement):
         ]
         return new_measurement
 
-    def shape(self, device, shots):
-        if not shots.has_partitioned_shots:
-            return ()
-        num_shot_elements = sum(s.copies for s in shots.shot_vector)
-        return tuple(() for _ in range(num_shot_elements))
+    def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple:
+        return ()
 
     def process_state(self, state: Sequence[complex], wire_order: Wires):
         state = qml.math.dm_from_state_vector(state)
diff --git a/pennylane/measurements/probs.py b/pennylane/measurements/probs.py
index 74e62bf12..5bc71b5ed 100644
--- a/pennylane/measurements/probs.py
+++ b/pennylane/measurements/probs.py
@@ -163,16 +163,9 @@ class ProbabilityMP(SampleMeasurement, StateMeasurement):
     def numeric_type(self):
         return float
 
-    def shape(self, device, shots):
-        num_shot_elements = (
-            sum(s.copies for s in shots.shot_vector) if shots.has_partitioned_shots else 1
-        )
-        len_wires = len(self.wires)
-        if len_wires == 0:
-            len_wires = len(device.wires) if device.wires else 0
-        dim = self._get_num_basis_states(len_wires, device)
-
-        return (dim,) if num_shot_elements == 1 else tuple((dim,) for _ in range(num_shot_elements))
+    def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple[int]:
+        len_wires = len(self.wires) if self.wires else num_device_wires
+        return (2**len_wires,)
 
     def process_samples(
         self,
diff --git a/pennylane/measurements/purity.py b/pennylane/measurements/purity.py
index e071f05dd..83e4166cb 100644
--- a/pennylane/measurements/purity.py
+++ b/pennylane/measurements/purity.py
@@ -85,11 +85,8 @@ class PurityMP(StateMeasurement):
     def numeric_type(self):
         return float
 
-    def shape(self, device, shots):
-        if not shots.has_partitioned_shots:
-            return ()
-        num_shot_elements = sum(s.copies for s in shots.shot_vector)
-        return tuple(() for _ in range(num_shot_elements))
+    def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple:
+        return ()
 
     def process_state(self, state: Sequence[complex], wire_order: Wires):
         wire_map = dict(zip(wire_order, list(range(len(wire_order)))))
diff --git a/pennylane/measurements/sample.py b/pennylane/measurements/sample.py
index b50c7f0a0..e341cbe75 100644
--- a/pennylane/measurements/sample.py
+++ b/pennylane/measurements/sample.py
@@ -220,7 +220,7 @@ class SampleMP(SampleMeasurement):
             return int
         return float
 
-    def shape(self, device, shots):
+    def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple:
         if not shots:
             raise MeasurementShapeError(
                 "Shots are required to obtain the shape of the measurement "
@@ -230,25 +230,13 @@ class SampleMP(SampleMeasurement):
             num_values_per_shot = 1  # one single eigenvalue
         else:
             # one value per wire
-            num_values_per_shot = len(self.wires) if len(self.wires) > 0 else len(device.wires)
-
-        def _single_int_shape(shot_val, num_values):
-            # singleton dimensions, whether in shot val or num_wires are squeezed away
-            inner_shape = []
-            if shot_val != 1:
-                inner_shape.append(shot_val)
-            if num_values != 1:
-                inner_shape.append(num_values)
-            return tuple(inner_shape)
-
-        if not shots.has_partitioned_shots:
-            return _single_int_shape(shots.total_shots, num_values_per_shot)
+            num_values_per_shot = len(self.wires) if len(self.wires) > 0 else num_device_wires
 
         shape = []
-        for s in shots.shot_vector:
-            for _ in range(s.copies):
-                shape.append(_single_int_shape(s.shots, num_values_per_shot))
-
+        if shots != 1:
+            shape.append(shots)
+        if num_values_per_shot != 1:
+            shape.append(num_values_per_shot)
         return tuple(shape)
 
     def process_samples(
diff --git a/pennylane/measurements/state.py b/pennylane/measurements/state.py
index b69934271..5b8baecf8 100644
--- a/pennylane/measurements/state.py
+++ b/pennylane/measurements/state.py
@@ -159,12 +159,9 @@ class StateMP(StateMeasurement):
     def numeric_type(self):
         return complex
 
-    def shape(self, device, shots):
-        num_shot_elements = (
-            sum(s.copies for s in shots.shot_vector) if shots.has_partitioned_shots else 1
-        )
-        dim = 2 ** len(self.wires) if self.wires else 2 ** len(device.wires)
-        return (dim,) if num_shot_elements == 1 else tuple((dim,) for _ in range(num_shot_elements))
+    def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple[int]:
+        num_wires = len(self.wires) if self.wires else num_device_wires
+        return (2**num_wires,)
 
     def process_state(self, state: Sequence[complex], wire_order: Wires):
         # pylint:disable=redefined-outer-name
@@ -232,17 +229,9 @@ class DensityMatrixMP(StateMP):
         shape = (2**n_wires, 2**n_wires)
         return shape, complex
 
-    def shape(self, device, shots):
-        num_shot_elements = (
-            sum(s.copies for s in shots.shot_vector) if shots.has_partitioned_shots else 1
-        )
-
+    def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple[int, int]:
         dim = 2 ** len(self.wires)
-        return (
-            (dim, dim)
-            if num_shot_elements == 1
-            else tuple((dim, dim) for _ in range(num_shot_elements))
-        )
+        return (dim, dim)
 
     def process_state(self, state: Sequence[complex], wire_order: Wires):
         # pylint:disable=redefined-outer-name
diff --git a/pennylane/measurements/var.py b/pennylane/measurements/var.py
index 21c9bf134..ab0448ffc 100644
--- a/pennylane/measurements/var.py
+++ b/pennylane/measurements/var.py
@@ -90,11 +90,8 @@ class VarianceMP(SampleMeasurement, StateMeasurement):
     def numeric_type(self):
         return float
 
-    def shape(self, device, shots):
-        if not shots.has_partitioned_shots:
-            return ()
-        num_shot_elements = sum(s.copies for s in shots.shot_vector)
-        return tuple(() for _ in range(num_shot_elements))
+    def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple:
+        return ()
 
     def process_samples(
         self,
diff --git a/pennylane/measurements/vn_entropy.py b/pennylane/measurements/vn_entropy.py
index 1589cee7b..79c1a875e 100644
--- a/pennylane/measurements/vn_entropy.py
+++ b/pennylane/measurements/vn_entropy.py
@@ -113,11 +113,8 @@ class VnEntropyMP(StateMeasurement):
     def numeric_type(self):
         return float
 
-    def shape(self, device, shots):
-        if not shots.has_partitioned_shots:
-            return ()
-        num_shot_elements = sum(s.copies for s in shots.shot_vector)
-        return tuple(() for _ in range(num_shot_elements))
+    def shape(self, shots: Optional[int] = None, num_device_wires: int = 0) -> tuple:
+        return ()
 
     def process_state(self, state: Sequence[complex], wire_order: Wires):
         state = qml.math.dm_from_state_vector(state)
diff --git a/pennylane/tape/qscript.py b/pennylane/tape/qscript.py
index c3b514d1c..9758e8fc1 100644
--- a/pennylane/tape/qscript.py
+++ b/pennylane/tape/qscript.py
@@ -773,7 +773,7 @@ class QuantumScript:
             dependent on the device used for execution.
 
         Args:
-            device (pennylane.Device): the device that will be used for the script execution
+            device (pennylane.devices.Device): the device that will be used for the script execution
 
         Returns:
             Union[tuple[int], tuple[tuple[int]]]: the output shape(s) of the quantum script result
@@ -791,27 +791,22 @@ class QuantumScript:
             >>> qs.shape(dev)
             ((4,), (), (4,))
         """
-        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(
-                "Parameter broadcasting when using a shot vector is not supported yet."
-            )
-
-        shapes = tuple(meas_process.shape(device, shots) for meas_process in self.measurements)
-
-        if self.batch_size is not None:
-            shapes = tuple((self.batch_size,) + shape for shape in shapes)
-
-        if len(shapes) == 1:
-            return shapes[0]
-
-        if shots.num_copies > 1:
-            # put the shot vector axis before the measurement axis
-            shapes = tuple(zip(*shapes))
-
-        return shapes
+        num_device_wires = len(device.wires) if device.wires else len(self.wires)
+
+        def get_shape(mp, _shots):
+            # depends on num_device_wires and self.batch_size from closure
+            standard_shape = mp.shape(shots=_shots, num_device_wires=num_device_wires)
+            if self.batch_size:
+                return (self.batch_size, *standard_shape)
+            return standard_shape
+
+        shape = []
+        for s in self.shots if self.shots else [None]:
+            shots_shape = tuple(get_shape(mp, s) for mp in self.measurements)
+            shots_shape = shots_shape[0] if len(shots_shape) == 1 else tuple(shots_shape)
+            shape.append(shots_shape)
+
+        return tuple(shape) if self.shots.has_partitioned_shots else shape[0]
 
     @property
     def numeric_type(self) -> Union[type, tuple[type, ...]]:
diff --git a/pennylane/workflow/interfaces/jax_jit.py b/pennylane/workflow/interfaces/jax_jit.py
index 552748748..f00028f07 100644
--- a/pennylane/workflow/interfaces/jax_jit.py
+++ b/pennylane/workflow/interfaces/jax_jit.py
@@ -58,9 +58,8 @@ def _to_jax(result: qml.typing.ResultBatch) -> qml.typing.ResultBatch:
         ResultBatch: a nested structure of tuples, and jax arrays
 
     """
-    # jax-jit not compatible with counts
-    # if isinstance(result, dict):
-    #     return result
+    if isinstance(result, dict):
+        return {key: jnp.array(value) for key, value in result.items()}
     if isinstance(result, (list, tuple)):
         return tuple(_to_jax(r) for r in result)
     return jnp.array(result)
@@ -86,23 +85,40 @@ def _jax_dtype(m_type):
     return jnp.dtype(m_type)
 
 
+def _get_counts_shape(mp: "qml.measurements.CountsMP", num_device_wires=0):
+    num_wires = len(mp.wires) if mp.wires else num_device_wires
+    outcome_counts = {}
+    binary_pattern = "{0:0" + str(num_wires) + "b}"
+    for outcome in range(2**num_wires):
+        outcome_binary = binary_pattern.format(outcome)
+        outcome_counts[outcome_binary] = jax.core.ShapedArray((), _jax_dtype(int))
+
+    return outcome_counts
+
+
 def _result_shape_dtype_struct(tape: "qml.tape.QuantumScript", device: "qml.Device"):
     """Auxiliary function for creating the shape and dtype object structure
     given a tape."""
 
-    shape = tape.shape(device)
-    if len(tape.measurements) == 1:
-        m_dtype = _jax_dtype(tape.measurements[0].numeric_type)
-        if tape.shots.has_partitioned_shots:
-            return tuple(jax.ShapeDtypeStruct(s, m_dtype) for s in shape)
-        return jax.ShapeDtypeStruct(tuple(shape), m_dtype)
-
-    tape_dtype = tuple(_jax_dtype(m.numeric_type) for m in tape.measurements)
-    if tape.shots.has_partitioned_shots:
-        return tuple(
-            tuple(jax.ShapeDtypeStruct(tuple(s), d) for s, d in zip(si, tape_dtype)) for si in shape
-        )
-    return tuple(jax.ShapeDtypeStruct(tuple(s), d) for s, d in zip(shape, tape_dtype))
+    num_device_wires = len(device.wires) if device.wires else len(tape.wires)
+
+    def struct(mp, shots):
+        # depends on num_device_wires and tape.batch_size from closure
+        if isinstance(mp, qml.measurements.CountsMP):
+            return _get_counts_shape(mp)
+        mp_shape = mp.shape(shots=shots, num_device_wires=num_device_wires)
+        if tape.batch_size:
+            mp_shape = (tape.batch_size, *mp_shape)
+        return jax.ShapeDtypeStruct(mp_shape, _jax_dtype(mp.numeric_type))
+
+    shape = []
+    for s in tape.shots if tape.shots else [None]:
+        shots_shape = tuple(struct(mp, s) for mp in tape.measurements)
+
+        shots_shape = shots_shape[0] if len(shots_shape) == 1 else tuple(shots_shape)
+        shape.append(shots_shape)
+
+    return tuple(shape) if tape.shots.has_partitioned_shots else shape[0]
 
 
 def _jac_shape_dtype_struct(tape: "qml.tape.QuantumScript", device: "qml.Device"):
@@ -234,11 +250,7 @@ def jax_jit_jvp_execute(tapes, execute_fn, jpc, device):
 
     """
 
-    if any(
-        m.return_type in (qml.measurements.Counts, qml.measurements.AllCounts)
-        for t in tapes
-        for m in t.measurements
-    ):
+    if any(m.return_type == qml.measurements.Counts for t in tapes for m in t.measurements):
         # Obtaining information about the shape of the Counts measurements is
         # not implemented and is required for the callback logic
         raise NotImplementedError("The JAX-JIT interface doesn't support qml.counts.")
@@ -264,11 +276,7 @@ def jax_jit_vjp_execute(tapes, execute_fn, jpc, device=None):
         the returned tuple corresponds in order to the provided tapes.
 
     """
-    if any(
-        m.return_type in (qml.measurements.Counts, qml.measurements.AllCounts)
-        for t in tapes
-        for m in t.measurements
-    ):
+    if any(m.return_type == qml.measurements.Counts for t in tapes for m in t.measurements):
         # Obtaining information about the shape of the Counts measurements is
         # not implemented and is required for the callback logic
         raise NotImplementedError("The JAX-JIT interface doesn't support qml.counts.")
diff --git a/pennylane/workflow/qnode.py b/pennylane/workflow/qnode.py
index 3b8b3c26f..d491e1f45 100644
--- a/pennylane/workflow/qnode.py
+++ b/pennylane/workflow/qnode.py
@@ -28,7 +28,7 @@ from cachetools import Cache
 import pennylane as qml
 from pennylane.debugging import pldb_device_manager
 from pennylane.logging import debug_logger
-from pennylane.measurements import CountsMP, MidMeasureMP
+from pennylane.measurements import MidMeasureMP
 from pennylane.tape import QuantumScript, QuantumTape
 from pennylane.transforms.core import TransformContainer, TransformDispatcher, TransformProgram
 
@@ -855,11 +855,6 @@ class QNode:
         params = self.tape.get_parameters(trainable_only=False)
         self.tape.trainable_params = qml.math.get_trainable_indices(params)
 
-        if any(isinstance(m, CountsMP) for m in self.tape.measurements) and any(
-            qml.math.is_abstract(a) for a in args
-        ):
-            raise qml.QuantumFunctionError("Can't JIT a quantum function that returns counts.")
-
         if isinstance(self._qfunc_output, qml.numpy.ndarray):
             measurement_processes = tuple(self.tape.measurements)
         elif not isinstance(self._qfunc_output, Sequence):


Test patch:
diff --git a/tests/devices/test_null_qubit.py b/tests/devices/test_null_qubit.py
index e206dae29..2af739553 100644
--- a/tests/devices/test_null_qubit.py
+++ b/tests/devices/test_null_qubit.py
@@ -466,13 +466,11 @@ class TestSampleMeasurements:
 
         dev = NullQubit()
         result = dev.execute(qs)
-        print(result)
 
-        assert (
-            result == tuple([{"00": s, "01": 0, "10": 0, "11": 0}] * 2 for s in qs.shots)
-            if all_outcomes
-            else tuple([{"00": s}] * 2 for s in qs.shots)
-        )
+        if all_outcomes:
+            assert result == tuple(({"00": s, "01": 0, "10": 0, "11": 0},) * 2 for s in qs.shots)
+        else:
+            assert result == tuple(({"00": s},) * 2 for s in qs.shots)
 
     @pytest.mark.parametrize("all_outcomes", [False, True])
     def test_counts_obs(self, all_outcomes):
@@ -500,11 +498,10 @@ class TestSampleMeasurements:
 
         dev = NullQubit()
         result = dev.execute(qs)
-        print(result)
         assert (
-            result == tuple([{-1: s, 1: 0}] * 2 for s in qs.shots)
+            result == tuple(({-1: s, 1: 0},) * 2 for s in qs.shots)
             if all_outcomes
-            else tuple([{-1: s}] * 2 for s in qs.shots)
+            else tuple(({-1: s},) * 2 for s in qs.shots)
         )
 
 
@@ -929,7 +926,6 @@ class TestDeviceDifferentiation:
         cotangents = [(0.456,), (0.789, 0.123)]
 
         actual_grad = dev.compute_vjp([single_meas, multi_meas], cotangents, self.ec)
-        print(actual_grad)
         assert actual_grad == ((0.0,), (0.0,))
 
         actual_val, actual_grad = dev.execute_and_compute_vjp(
@@ -956,7 +952,6 @@ class TestDeviceDifferentiation:
 
         assert new_ec.use_device_gradient
         assert new_ec.grad_on_execution
-        print(actual_grad)
         assert actual_grad == ((0.0,), (0.0,))
 
 
diff --git a/tests/interfaces/test_jax_jit.py b/tests/interfaces/test_jax_jit.py
index 70e5c41df..68ac56f2d 100644
--- a/tests/interfaces/test_jax_jit.py
+++ b/tests/interfaces/test_jax_jit.py
@@ -884,6 +884,22 @@ class TestVectorValuedJIT:
             assert jax.numpy.allclose(r, e, atol=1e-7)
 
 
+def test_jit_allcounts():
+    """Test jitting with counts with all_outcomes == True."""
+
+    tape = qml.tape.QuantumScript(
+        [qml.RX(0, 0)], [qml.counts(wires=(0, 1), all_outcomes=True)], shots=50
+    )
+    device = qml.device("default.qubit")
+
+    res = jax.jit(qml.execute, static_argnums=(1, 2))((tape,), device, qml.gradients.param_shift)[0]
+
+    assert set(res.keys()) == {"00", "01", "10", "11"}
+    assert qml.math.allclose(res["00"], 50)
+    for val in ["01", "10", "11"]:
+        assert qml.math.allclose(res[val], 0)
+
+
 @pytest.mark.xfail(reason="Need to figure out how to handle this case in a less ambiguous manner")
 def test_diff_method_None_jit():
     """Test that jitted execution works when `gradient_fn=None`."""
diff --git a/tests/measurements/legacy/test_classical_shadow_legacy.py b/tests/measurements/legacy/test_classical_shadow_legacy.py
index 053c9c797..747028f6a 100644
--- a/tests/measurements/legacy/test_classical_shadow_legacy.py
+++ b/tests/measurements/legacy/test_classical_shadow_legacy.py
@@ -18,7 +18,6 @@ import pytest
 
 import pennylane as qml
 from pennylane import numpy as np
-from pennylane.measurements import Shots
 
 # pylint: disable=dangerous-default-value, too-many-arguments
 
@@ -73,20 +72,6 @@ class TestClassicalShadow:
     shots_list = [1, 100]
     seed_recipes_list = [None, 74]  # random seed
 
-    @pytest.mark.parametrize("shots", shots_list)
-    @pytest.mark.parametrize("seed", seed_recipes_list)
-    def test_measurement_process_shape(self, wires, shots, seed):
-        """Test that the shape of the MeasurementProcess instance is correct"""
-        dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
-        shots_obj = Shots(shots)
-        res = qml.classical_shadow(wires=range(wires), seed=seed)
-        assert res.shape(dev, shots_obj) == (2, shots, wires)
-
-        # test an error is raised when device is None
-        msg = "Shots must be specified to obtain the shape of a classical shadow measurement"
-        with pytest.raises(qml.measurements.MeasurementShapeError, match=msg):
-            res.shape(dev, Shots(None))
-
     def test_shape_matches(self, wires):
         """Test that the shape of the MeasurementProcess matches the shape
         of the tape execution"""
@@ -96,9 +81,7 @@ class TestClassicalShadow:
         circuit.construct((), {})
 
         res = qml.execute([circuit.tape], circuit.device, None)[0]
-        expected_shape = qml.classical_shadow(wires=range(wires)).shape(
-            circuit.device, Shots(shots)
-        )
+        expected_shape = qml.classical_shadow(wires=range(wires)).shape(shots, wires)
 
         assert res.shape == expected_shape
 
@@ -240,14 +223,6 @@ def qft_circuit(wires, shots=10000, interface="autograd"):
 
 @pytest.mark.autograd
 class TestExpvalMeasurement:
-    @pytest.mark.parametrize("wires", [1, 2])
-    @pytest.mark.parametrize("shots", [1, 10])
-    def test_measurement_process_shape(self, wires, shots):
-        """Test that the shape of the MeasurementProcess instance is correct"""
-        dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
-        H = qml.PauliZ(0)
-        res = qml.shadow_expval(H)
-        assert len(res.shape(dev, Shots(shots))) == 0
 
     def test_shape_matches(self):
         """Test that the shape of the MeasurementProcess matches the shape
@@ -260,7 +235,7 @@ class TestExpvalMeasurement:
         circuit.construct((H,), {})
 
         res = qml.execute([circuit.tape], circuit.device, None)[0]
-        expected_shape = qml.shadow_expval(H).shape(circuit.device, Shots(shots))
+        expected_shape = qml.shadow_expval(H).shape(shots, len(circuit.device.wires))
 
         assert res.shape == expected_shape
 
diff --git a/tests/measurements/legacy/test_expval_legacy.py b/tests/measurements/legacy/test_expval_legacy.py
index 6118b0ac0..d30f11dcd 100644
--- a/tests/measurements/legacy/test_expval_legacy.py
+++ b/tests/measurements/legacy/test_expval_legacy.py
@@ -17,7 +17,7 @@ import pytest
 
 import pennylane as qml
 from pennylane.devices.qubit.measure import flatten_state
-from pennylane.measurements import Expectation, Shots
+from pennylane.measurements import Expectation
 
 
 # TODO: Remove this when new CustomMP are the default
@@ -174,30 +174,6 @@ class TestExpval:
                 )
             custom_measurement_process(new_dev, spy)
 
-    @pytest.mark.parametrize(
-        "obs",
-        [qml.PauliZ(0), qml.Hermitian(np.diag([1, 2]), 0), qml.Hermitian(np.diag([1.0, 2.0]), 0)],
-    )
-    def test_shape(self, obs):
-        """Test that the shape is correct."""
-        dev = qml.device("default.qubit.legacy", wires=1)
-
-        res = qml.expval(obs)
-        # pylint: disable=use-implicit-booleaness-not-comparison
-        assert res.shape(dev, Shots(None)) == ()
-        assert res.shape(dev, Shots(100)) == ()
-
-    @pytest.mark.parametrize(
-        "obs",
-        [qml.PauliZ(0), qml.Hermitian(np.diag([1, 2]), 0), qml.Hermitian(np.diag([1.0, 2.0]), 0)],
-    )
-    def test_shape_shot_vector(self, obs):
-        """Test that the shape is correct with the shot vector too."""
-        res = qml.expval(obs)
-        shot_vector = (1, 2, 3)
-        dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vector)
-        assert res.shape(dev, Shots(shot_vector)) == ((), (), ())
-
     @pytest.mark.parametrize("state", [np.array([0, 0, 0]), np.array([1, 0, 0, 0, 0, 0, 0, 0])])
     @pytest.mark.parametrize("shots", [None, 1000, [1000, 10000]])
     def test_projector_expval(self, state, shots, mocker):
diff --git a/tests/measurements/legacy/test_mutual_info_legacy.py b/tests/measurements/legacy/test_mutual_info_legacy.py
index 023c79205..fee7b9f07 100644
--- a/tests/measurements/legacy/test_mutual_info_legacy.py
+++ b/tests/measurements/legacy/test_mutual_info_legacy.py
@@ -16,18 +16,9 @@ import numpy as np
 import pytest
 
 import pennylane as qml
-from pennylane.measurements import Shots
 from pennylane.workflow import INTERFACE_MAP
 
 
-@pytest.mark.parametrize("shots, shape", [(None, ()), (10, ()), ([1, 10], ((), ()))])
-def test_shape(shots, shape):
-    """Test that the shape is correct."""
-    dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
-    res = qml.mutual_info(wires0=[0], wires1=[1])
-    assert res.shape(dev, Shots(shots)) == shape
-
-
 class TestIntegration:
     """Tests for the mutual information functions"""
 
diff --git a/tests/measurements/legacy/test_probs_legacy.py b/tests/measurements/legacy/test_probs_legacy.py
index bbc326e77..bbe3284c7 100644
--- a/tests/measurements/legacy/test_probs_legacy.py
+++ b/tests/measurements/legacy/test_probs_legacy.py
@@ -18,7 +18,7 @@ import pytest
 import pennylane as qml
 from pennylane import numpy as pnp
 from pennylane.devices.qubit.measure import flatten_state
-from pennylane.measurements import ProbabilityMP, Shots
+from pennylane.measurements import ProbabilityMP
 
 
 # TODO: Remove this when new CustomMP are the default
@@ -79,26 +79,6 @@ class TestProbs:
 
         assert isinstance(circuit.tape[0], ProbabilityMP)
 
-    @pytest.mark.parametrize("wires", [[0], [2, 1], ["a", "c", 3]])
-    @pytest.mark.parametrize("shots", [None, 10])
-    def test_shape(self, wires, shots):
-        """Test that the shape is correct."""
-        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
-        res = qml.probs(wires=wires)
-        assert res.shape(dev, Shots(shots)) == (2 ** len(wires),)
-
-    @pytest.mark.parametrize("wires", [[0], [2, 1], ["a", "c", 3]])
-    def test_shape_shot_vector(self, wires):
-        """Test that the shape is correct with the shot vector too."""
-        res = qml.probs(wires=wires)
-        shot_vector = (1, 2, 3)
-        dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vector)
-        assert res.shape(dev, Shots(shot_vector)) == (
-            (2 ** len(wires),),
-            (2 ** len(wires),),
-            (2 ** len(wires),),
-        )
-
     @pytest.mark.parametrize("shots", [None, 100])
     def test_probs_no_arguments(self, shots):
         """Test that using ``qml.probs`` with no arguments returns the probabilities of all wires."""
diff --git a/tests/measurements/legacy/test_purity_measurement_legacy.py b/tests/measurements/legacy/test_purity_measurement_legacy.py
index 752b07b5f..4ff7a8508 100644
--- a/tests/measurements/legacy/test_purity_measurement_legacy.py
+++ b/tests/measurements/legacy/test_purity_measurement_legacy.py
@@ -17,7 +17,6 @@ import numpy as np
 import pytest
 
 import pennylane as qml
-from pennylane.measurements import Shots
 
 # pylint: disable=too-many-arguments
 
@@ -53,14 +52,6 @@ def expected_purity_grad_ising_xx(param):
     return grad_expected_purity
 
 
-@pytest.mark.parametrize("shots, shape", [(None, ()), (10, ()), ((1, 10), ((), ()))])
-def test_shape_new(shots, shape):
-    """Test the ``shape_new`` method."""
-    meas = qml.purity(wires=0)
-    dev = qml.device("default.qubit.legacy", wires=1, shots=shots)
-    assert meas.shape(dev, Shots(shots)) == shape
-
-
 class TestPurityIntegration:
     """Test the purity meausrement with qnodes and devices."""
 
diff --git a/tests/measurements/legacy/test_sample_legacy.py b/tests/measurements/legacy/test_sample_legacy.py
index b1bc69bb7..9f836e776 100644
--- a/tests/measurements/legacy/test_sample_legacy.py
+++ b/tests/measurements/legacy/test_sample_legacy.py
@@ -16,7 +16,7 @@ import numpy as np
 import pytest
 
 import pennylane as qml
-from pennylane.measurements import MeasurementShapeError, MeasurementValue, Sample, Shots
+from pennylane.measurements import MeasurementValue, Sample
 from pennylane.operation import Operator
 
 # pylint: disable=protected-access, no-member
@@ -62,10 +62,10 @@ class TestSample:
         output = circuit()
 
         assert len(output) == 2
-        assert circuit._qfunc_output[0].shape(dev, Shots(n_sample)) == (
+        assert circuit._qfunc_output[0].shape(n_sample, 2) == (
             (n_sample,) if not n_sample == 1 else ()
         )
-        assert circuit._qfunc_output[1].shape(dev, Shots(n_sample)) == (
+        assert circuit._qfunc_output[1].shape(n_sample, 2) == (
             (n_sample,) if not n_sample == 1 else ()
         )
 
@@ -89,7 +89,7 @@ class TestSample:
 
         assert len(result) == 3
         assert np.array_equal(result[0].shape, (n_sample,))
-        assert circuit._qfunc_output[0].shape(dev, Shots(n_sample)) == (n_sample,)
+        assert circuit._qfunc_output[0].shape(n_sample, 3) == (n_sample,)
         assert isinstance(result[1], np.ndarray)
         assert isinstance(result[2], np.ndarray)
 
@@ -198,7 +198,7 @@ class TestSample:
 
         assert isinstance(result, np.ndarray)
         assert np.array_equal(result.shape, (n_sample,))
-        assert circuit._qfunc_output.shape(dev, Shots(n_sample)) == (n_sample,)
+        assert circuit._qfunc_output.shape(n_sample, 1) == (n_sample,)
 
         custom_measurement_process(dev, spy)
 
@@ -216,9 +216,9 @@ class TestSample:
 
         result = circuit()
 
-        assert circuit._qfunc_output[0].shape(dev, Shots(n_sample)) == (n_sample,)
-        assert circuit._qfunc_output[1].shape(dev, Shots(n_sample)) == (n_sample,)
-        assert circuit._qfunc_output[2].shape(dev, Shots(n_sample)) == (n_sample,)
+        assert circuit._qfunc_output[0].shape(n_sample, 3) == (n_sample,)
+        assert circuit._qfunc_output[1].shape(n_sample, 3) == (n_sample,)
+        assert circuit._qfunc_output[2].shape(n_sample, 3) == (n_sample,)
 
         # If all the dimensions are equal the result will end up to be a proper rectangular array
         assert isinstance(result, tuple)
@@ -353,58 +353,6 @@ class TestSample:
 
         custom_measurement_process(dev, spy)
 
-    def test_shape_no_shots_error(self):
-        """Test that the appropriate error is raised with no shots are specified"""
-        dev = qml.device("default.qubit.legacy", wires=2, shots=None)
-        shots = Shots(None)
-        mp = qml.sample()
-
-        with pytest.raises(
-            MeasurementShapeError, match="Shots are required to obtain the shape of the measurement"
-        ):
-            _ = mp.shape(dev, shots)
-
-    @pytest.mark.parametrize(
-        "obs",
-        [
-            None,
-            qml.PauliZ(0),
-            qml.Hermitian(np.diag([1, 2]), 0),
-            qml.Hermitian(np.diag([1.0, 2.0]), 0),
-        ],
-    )
-    def test_shape(self, obs):
-        """Test that the shape is correct."""
-        shots = 10
-        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
-        res = qml.sample(obs) if obs is not None else qml.sample()
-        expected = (shots,) if obs is not None else (shots, 3)
-        assert res.shape(dev, Shots(shots)) == expected
-
-    @pytest.mark.parametrize("n_samples", (1, 10))
-    def test_shape_wires(self, n_samples):
-        """Test that the shape is correct when wires are provided."""
-        dev = qml.device("default.qubit.legacy", wires=3, shots=n_samples)
-        mp = qml.sample(wires=(0, 1))
-        assert mp.shape(dev, Shots(n_samples)) == (n_samples, 2) if n_samples != 1 else (2,)
-
-    @pytest.mark.parametrize(
-        "obs",
-        [
-            None,
-            qml.PauliZ(0),
-            qml.Hermitian(np.diag([1, 2]), 0),
-            qml.Hermitian(np.diag([1.0, 2.0]), 0),
-        ],
-    )
-    def test_shape_shot_vector(self, obs):
-        """Test that the shape is correct with the shot vector too."""
-        shot_vector = (1, 2, 3)
-        dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vector)
-        res = qml.sample(obs) if obs is not None else qml.sample()
-        expected = ((), (2,), (3,)) if obs is not None else ((3,), (2, 3), (3, 3))
-        assert res.shape(dev, Shots(shot_vector)) == expected
-
     def test_shape_shot_vector_obs(self):
         """Test that the shape is correct with the shot vector and a observable too."""
         shot_vec = (2, 2)
@@ -459,6 +407,4 @@ def test_jitting_with_sampling_on_subset_of_wires(samples):
 
     expected = (2,) if samples == 1 else (samples, 2)
     assert results.shape == expected
-    assert (
-        circuit._qfunc_output.shape(dev, Shots(samples)) == (samples, 2) if samples != 1 else (2,)
-    )
+    assert circuit._qfunc_output.shape(samples, 3) == (samples, 2) if samples != 1 else (2,)
diff --git a/tests/measurements/legacy/test_state_legacy.py b/tests/measurements/legacy/test_state_legacy.py
index 83e7cba2b..6e9cadcae 100644
--- a/tests/measurements/legacy/test_state_legacy.py
+++ b/tests/measurements/legacy/test_state_legacy.py
@@ -18,7 +18,7 @@ import pytest
 import pennylane as qml
 from pennylane import numpy as pnp
 from pennylane.devices import DefaultQubitLegacy
-from pennylane.measurements import Shots, State, density_matrix, expval, state
+from pennylane.measurements import State, density_matrix, expval, state
 
 
 class TestState:
@@ -324,20 +324,6 @@ class TestState:
 
         assert np.allclose(state_expected, state_val)
 
-    @pytest.mark.parametrize("shots", [None, 1, 10])
-    def test_shape(self, shots):
-        """Test that the shape is correct for qml.state."""
-        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
-        res = qml.state()
-        assert res.shape(dev, Shots(shots)) == (2**3,)
-
-    @pytest.mark.parametrize("s_vec", [(3, 2, 1), (1, 5, 10), (3, 1, 20)])
-    def test_shape_shot_vector(self, s_vec):
-        """Test that the shape is correct for qml.state with the shot vector too."""
-        dev = qml.device("default.qubit.legacy", wires=3, shots=s_vec)
-        res = qml.state()
-        assert res.shape(dev, Shots(s_vec)) == ((2**3,), (2**3,), (2**3,))
-
 
 class TestDensityMatrix:
     """Tests for the density matrix function"""
@@ -373,7 +359,7 @@ class TestDensityMatrix:
             return density_matrix(0)
 
         func()
-        obs = func.qtape.observables
+        obs = func.qtape.measurements
         assert len(obs) == 1
         assert obs[0].return_type is State
 
@@ -783,21 +769,3 @@ class TestDensityMatrix:
             ),
             density,
         )
-
-    @pytest.mark.parametrize("shots", [None, 1, 10])
-    def test_shape(self, shots):
-        """Test that the shape is correct for qml.density_matrix."""
-        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
-        res = qml.density_matrix(wires=[0, 1])
-        assert res.shape(dev, Shots(shots)) == (2**2, 2**2)
-
-    @pytest.mark.parametrize("s_vec", [(3, 2, 1), (1, 5, 10), (3, 1, 20)])
-    def test_shape_shot_vector(self, s_vec):
-        """Test that the shape is correct for qml.density_matrix with the shot vector too."""
-        dev = qml.device("default.qubit.legacy", wires=3, shots=s_vec)
-        res = qml.density_matrix(wires=[0, 1])
-        assert res.shape(dev, Shots(s_vec)) == (
-            (2**2, 2**2),
-            (2**2, 2**2),
-            (2**2, 2**2),
-        )
diff --git a/tests/measurements/legacy/test_var_legacy.py b/tests/measurements/legacy/test_var_legacy.py
index 70000bcd6..cea386b22 100644
--- a/tests/measurements/legacy/test_var_legacy.py
+++ b/tests/measurements/legacy/test_var_legacy.py
@@ -17,7 +17,7 @@ import pytest
 
 import pennylane as qml
 from pennylane.devices.qubit.measure import flatten_state
-from pennylane.measurements import Shots, Variance
+from pennylane.measurements import Variance
 
 
 # TODO: Remove this when new CustomMP are the default
@@ -169,29 +169,6 @@ class TestVar:
                 new_dev.target_device._samples = new_dev.generate_samples()
             custom_measurement_process(new_dev, spy)
 
-    @pytest.mark.parametrize(
-        "obs",
-        [qml.PauliZ(0), qml.Hermitian(np.diag([1, 2]), 0), qml.Hermitian(np.diag([1.0, 2.0]), 0)],
-    )
-    def test_shape(self, obs):
-        """Test that the shape is correct."""
-        dev = qml.device("default.qubit.legacy", wires=1)
-        res = qml.var(obs)
-        # pylint: disable=use-implicit-booleaness-not-comparison
-        assert res.shape(dev, Shots(None)) == ()
-        assert res.shape(dev, Shots(100)) == ()
-
-    @pytest.mark.parametrize(
-        "obs",
-        [qml.PauliZ(0), qml.Hermitian(np.diag([1, 2]), 0), qml.Hermitian(np.diag([1.0, 2.0]), 0)],
-    )
-    def test_shape_shot_vector(self, obs):
-        """Test that the shape is correct with the shot vector too."""
-        res = qml.var(obs)
-        shot_vector = (1, 2, 3)
-        dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vector)
-        assert res.shape(dev, Shots(shot_vector)) == ((), (), ())
-
     @pytest.mark.parametrize("state", [np.array([0, 0, 0]), np.array([1, 0, 0, 0, 0, 0, 0, 0])])
     @pytest.mark.parametrize("shots", [None, 1000, [1000, 10000]])
     def test_projector_var(self, state, shots, mocker):
diff --git a/tests/measurements/legacy/test_vn_entropy_legacy.py b/tests/measurements/legacy/test_vn_entropy_legacy.py
index 9a4dbf2f7..6a9f53b3b 100644
--- a/tests/measurements/legacy/test_vn_entropy_legacy.py
+++ b/tests/measurements/legacy/test_vn_entropy_legacy.py
@@ -16,7 +16,6 @@ import numpy as np
 import pytest
 
 import pennylane as qml
-from pennylane.measurements import Shots
 from pennylane.measurements.vn_entropy import VnEntropyMP
 from pennylane.workflow import INTERFACE_MAP
 
@@ -101,14 +100,6 @@ class TestInitialization:
 
         assert isinstance(circuit.tape[0], VnEntropyMP)
 
-    @pytest.mark.parametrize("shots, shape", [(None, ()), (10, ()), ((1, 10), ((), ()))])
-    def test_shape(self, shots, shape):
-        """Test the ``shape`` method."""
-        meas = qml.vn_entropy(wires=0)
-        dev = qml.device("default.qubit.legacy", wires=1, shots=shots)
-
-        assert meas.shape(dev, Shots(shots)) == shape
-
 
 class TestIntegration:
     """Integration tests for the vn_entropy measurement function."""
diff --git a/tests/measurements/test_classical_shadow.py b/tests/measurements/test_classical_shadow.py
index 04d9f8e80..a900a1313 100644
--- a/tests/measurements/test_classical_shadow.py
+++ b/tests/measurements/test_classical_shadow.py
@@ -20,7 +20,7 @@ import pytest
 
 import pennylane as qml
 from pennylane import numpy as np
-from pennylane.measurements import ClassicalShadowMP, Shots
+from pennylane.measurements import ClassicalShadowMP
 from pennylane.measurements.classical_shadow import ShadowExpvalMP
 
 # pylint: disable=dangerous-default-value, too-many-arguments
@@ -239,15 +239,13 @@ class TestClassicalShadow:
     @pytest.mark.parametrize("seed", seed_recipes_list)
     def test_measurement_process_shape(self, wires, shots, seed):
         """Test that the shape of the MeasurementProcess instance is correct"""
-        dev = qml.device("default.qubit", wires=wires, shots=shots)
-        shots_obj = Shots(shots)
         res = qml.classical_shadow(wires=range(wires), seed=seed)
-        assert res.shape(dev, shots_obj) == (2, shots, wires)
+        assert res.shape(shots, wires) == (2, shots, wires)
 
         # test an error is raised when device is None
         msg = "Shots must be specified to obtain the shape of a classical shadow measurement"
         with pytest.raises(qml.measurements.MeasurementShapeError, match=msg):
-            res.shape(dev, Shots(None))
+            res.shape(None, wires)
 
     def test_shape_matches(self, wires):
         """Test that the shape of the MeasurementProcess matches the shape
@@ -258,9 +256,7 @@ class TestClassicalShadow:
         circuit.construct((), {})
 
         res = qml.execute([circuit.tape], circuit.device, None)[0]
-        expected_shape = qml.classical_shadow(wires=range(wires)).shape(
-            circuit.device, Shots(shots)
-        )
+        expected_shape = qml.classical_shadow(wires=range(wires)).shape(shots, wires)
 
         assert res.shape == expected_shape
 
@@ -422,10 +418,9 @@ class TestExpvalMeasurement:
     @pytest.mark.parametrize("shots", [1, 10])
     def test_measurement_process_shape(self, wires, shots):
         """Test that the shape of the MeasurementProcess instance is correct"""
-        dev = qml.device("default.qubit", wires=wires, shots=shots)
         H = qml.PauliZ(0)
         res = qml.shadow_expval(H)
-        assert len(res.shape(dev, Shots(shots))) == 0
+        assert len(res.shape(shots, wires)) == 0
 
     def test_shape_matches(self):
         """Test that the shape of the MeasurementProcess matches the shape
@@ -438,7 +433,7 @@ class TestExpvalMeasurement:
         circuit.construct((H,), {})
 
         res = qml.execute([circuit.tape], circuit.device, None)[0]
-        expected_shape = qml.shadow_expval(H).shape(circuit.device, Shots(shots))
+        expected_shape = qml.shadow_expval(H).shape(shots, wires)
 
         assert res.shape == expected_shape
 
diff --git a/tests/measurements/test_expval.py b/tests/measurements/test_expval.py
index 14dcbb4ba..b1a5e7178 100644
--- a/tests/measurements/test_expval.py
+++ b/tests/measurements/test_expval.py
@@ -18,7 +18,7 @@ import numpy as np
 import pytest
 
 import pennylane as qml
-from pennylane.measurements import Expectation, Shots
+from pennylane.measurements import Expectation
 from pennylane.measurements.expval import ExpectationMP
 
 
@@ -180,23 +180,11 @@ class TestExpval:
     )
     def test_shape(self, obs):
         """Test that the shape is correct."""
-        dev = qml.device("default.qubit", wires=1)
 
         res = qml.expval(obs)
         # pylint: disable=use-implicit-booleaness-not-comparison
-        assert res.shape(dev, Shots(None)) == ()
-        assert res.shape(dev, Shots(100)) == ()
-
-    @pytest.mark.parametrize(
-        "obs",
-        [qml.PauliZ(0), qml.Hermitian(np.diag([1, 2]), 0), qml.Hermitian(np.diag([1.0, 2.0]), 0)],
-    )
-    def test_shape_shot_vector(self, obs):
-        """Test that the shape is correct with the shot vector too."""
-        res = qml.expval(obs)
-        shot_vector = (1, 2, 3)
-        dev = qml.device("default.qubit", wires=3, shots=shot_vector)
-        assert res.shape(dev, Shots(shot_vector)) == ((), (), ())
+        assert res.shape(None, 1) == ()
+        assert res.shape(100, 1) == ()
 
     @pytest.mark.parametrize("state", [np.array([0, 0, 0]), np.array([1, 0, 0, 0, 0, 0, 0, 0])])
     @pytest.mark.parametrize("shots", [None, 1000, [1000, 1111]])
diff --git a/tests/measurements/test_measurements.py b/tests/measurements/test_measurements.py
index 20c789fe2..3019da7dd 100644
--- a/tests/measurements/test_measurements.py
+++ b/tests/measurements/test_measurements.py
@@ -687,3 +687,50 @@ class TestMeasurementTransform:
             return CountTapesMP(wires=[0])
 
         assert circuit() == 1
+
+
+class TestMeasurementProcess:
+    """Tests for the shape and numeric type of a measurement process"""
+
+    measurements_no_shots = [
+        (qml.expval(qml.PauliZ(0)), ()),
+        (qml.var(qml.PauliZ(0)), ()),
+        (qml.probs(wires=[0, 1]), (4,)),
+        (qml.state(), (8,)),
+        (qml.density_matrix(wires=[0, 1]), (4, 4)),
+        (qml.mutual_info(wires0=[0], wires1=[1]), ()),
+        (qml.vn_entropy(wires=[0, 1]), ()),
+    ]
+
+    measurements_finite_shots = [
+        (qml.expval(qml.PauliZ(0)), ()),
+        (qml.var(qml.PauliZ(0)), ()),
+        (qml.probs(wires=[0, 1]), (4,)),
+        (qml.state(), (8,)),
+        (qml.density_matrix(wires=[0, 1]), (4, 4)),
+        (qml.sample(qml.PauliZ(0)), (10,)),
+        (qml.sample(), (10, 3)),
+        (qml.mutual_info(wires0=0, wires1=1), ()),
+        (qml.vn_entropy(wires=[0, 1]), ()),
+    ]
+
+    @pytest.mark.parametrize("measurement, expected_shape", measurements_no_shots)
+    def test_output_shapes_no_shots(self, measurement, expected_shape):
+        """Test that the output shape of the measurement process is expected
+        when shots=None"""
+
+        assert measurement.shape(shots=None, num_device_wires=3) == expected_shape
+
+    @pytest.mark.parametrize("measurement, expected_shape", measurements_finite_shots)
+    def test_output_shapes_finite_shots(self, measurement, expected_shape):
+        """Test that the output shape of the measurement process is expected
+        when shots is finite"""
+        assert measurement.shape(shots=10, num_device_wires=3) == expected_shape
+
+    def test_undefined_shape_error(self):
+        """Test that an error is raised for a measurement with an undefined shape"""
+        measurement = qml.counts(wires=[0, 1])
+        msg = "The shape of the measurement CountsMP is not defined"
+
+        with pytest.raises(qml.QuantumFunctionError, match=msg):
+            measurement.shape(shots=None, num_device_wires=2)
diff --git a/tests/measurements/test_mutual_info.py b/tests/measurements/test_mutual_info.py
index c457afd6b..a47284f74 100644
--- a/tests/measurements/test_mutual_info.py
+++ b/tests/measurements/test_mutual_info.py
@@ -18,7 +18,7 @@ import numpy as np
 import pytest
 
 import pennylane as qml
-from pennylane.measurements import MutualInfo, Shots
+from pennylane.measurements import MutualInfo
 from pennylane.measurements.mutual_info import MutualInfoMP
 from pennylane.wires import Wires
 
@@ -35,12 +35,11 @@ class TestMutualInfoUnitTests:
         assert q.queue[0] is m
         assert isinstance(q.queue[0], MutualInfoMP)
 
-    @pytest.mark.parametrize("shots, shape", [(None, ()), (10, ()), ([1, 10], ((), ()))])
+    @pytest.mark.parametrize("shots, shape", [(None, ()), (10, ())])
     def test_shape(self, shots, shape):
         """Test that the shape is correct."""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
         res = qml.mutual_info(wires0=[0], wires1=[1])
-        assert res.shape(dev, Shots(shots)) == shape
+        assert res.shape(shots, 3) == shape
 
     def test_properties(self):
         """Test that the properties are correct."""
diff --git a/tests/measurements/test_probs.py b/tests/measurements/test_probs.py
index 83e76fca5..48d40d5d5 100644
--- a/tests/measurements/test_probs.py
+++ b/tests/measurements/test_probs.py
@@ -19,7 +19,7 @@ import pytest
 
 import pennylane as qml
 from pennylane import numpy as pnp
-from pennylane.measurements import MeasurementProcess, Probability, ProbabilityMP, Shots
+from pennylane.measurements import MeasurementProcess, Probability, ProbabilityMP
 from pennylane.queuing import AnnotatedQueue
 
 
@@ -63,31 +63,16 @@ class TestProbs:
     @pytest.mark.parametrize("shots", [None, 10])
     def test_shape(self, wires, shots):
         """Test that the shape is correct."""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
         res = qml.probs(wires=wires)
-        assert res.shape(dev, Shots(shots)) == (2 ** len(wires),)
+        assert res.shape(shots, 3) == (2 ** len(wires),)
 
     def test_shape_empty_wires(self):
         """Test that shape works when probs is broadcasted onto all available wires."""
-        dev = qml.device("default.qubit", wires=(1, 2, 3))
         res = qml.probs()
-        assert res.shape(dev, Shots(None)) == (8,)
+        assert res.shape(None, 3) == (8,)
 
-        dev2 = qml.device("default.qubit")
         res = qml.probs()
-        assert res.shape(dev2, Shots(None)) == (1,)
-
-    @pytest.mark.parametrize("wires", [[0], [2, 1], ["a", "c", 3]])
-    def test_shape_shot_vector(self, wires):
-        """Test that the shape is correct with the shot vector too."""
-        res = qml.probs(wires=wires)
-        shot_vector = (1, 2, 3)
-        dev = qml.device("default.qubit", wires=3, shots=shot_vector)
-        assert res.shape(dev, Shots(shot_vector)) == (
-            (2 ** len(wires),),
-            (2 ** len(wires),),
-            (2 ** len(wires),),
-        )
+        assert res.shape(None, 0) == (1,)
 
     @pytest.mark.parametrize("wires", [[0], [0, 1], [1, 0, 2]])
     def test_annotating_probs(self, wires):
diff --git a/tests/measurements/test_purity_measurement.py b/tests/measurements/test_purity_measurement.py
index 2ed4a9ea5..819051023 100644
--- a/tests/measurements/test_purity_measurement.py
+++ b/tests/measurements/test_purity_measurement.py
@@ -17,7 +17,7 @@ import numpy as np
 import pytest
 
 import pennylane as qml
-from pennylane.measurements import PurityMP, Shots
+from pennylane.measurements import PurityMP
 
 # pylint: disable=too-many-arguments
 
@@ -66,12 +66,11 @@ class TestPurityUnitTest:
         m = PurityMP(wires=qml.wires.Wires(0))
         assert m.numeric_type is float
 
-    @pytest.mark.parametrize("shots, shape", [(None, ()), (10, ()), ((1, 10), ((), ()))])
+    @pytest.mark.parametrize("shots, shape", [(None, ()), (10, ())])
     def test_shape_new(self, shots, shape):
         """Test the ``shape_new`` method."""
         meas = qml.purity(wires=0)
-        dev = qml.device("default.qubit", wires=1, shots=shots)
-        assert meas.shape(dev, Shots(shots)) == shape
+        assert meas.shape(shots, 1) == shape
 
 
 class TestPurityIntegration:
diff --git a/tests/measurements/test_sample.py b/tests/measurements/test_sample.py
index c54b78f62..fb181123f 100644
--- a/tests/measurements/test_sample.py
+++ b/tests/measurements/test_sample.py
@@ -16,7 +16,7 @@ import numpy as np
 import pytest
 
 import pennylane as qml
-from pennylane.measurements import MeasurementShapeError, Sample, Shots
+from pennylane.measurements import MeasurementShapeError, Sample
 from pennylane.operation import EigvalsUndefinedError, Operator
 
 # pylint: disable=protected-access, no-member, too-many-public-methods
@@ -39,10 +39,10 @@ class TestSample:
         output = circuit()
 
         assert len(output) == 2
-        assert circuit._qfunc_output[0].shape(dev, Shots(n_sample)) == (
+        assert circuit._qfunc_output[0].shape(shots=n_sample, num_device_wires=2) == (
             (n_sample,) if not n_sample == 1 else ()
         )
-        assert circuit._qfunc_output[1].shape(dev, Shots(n_sample)) == (
+        assert circuit._qfunc_output[1].shape(shots=n_sample, num_device_wires=2) == (
             (n_sample,) if not n_sample == 1 else ()
         )
 
@@ -62,7 +62,7 @@ class TestSample:
 
         assert len(result) == 3
         assert np.array_equal(result[0].shape, (n_sample,))
-        assert circuit._qfunc_output[0].shape(dev, Shots(n_sample)) == (n_sample,)
+        assert circuit._qfunc_output[0].shape(shots=n_sample, num_device_wires=3) == (n_sample,)
         assert isinstance(result[1], np.float64)
         assert isinstance(result[2], np.float64)
 
@@ -81,7 +81,7 @@ class TestSample:
 
         assert isinstance(result, np.ndarray)
         assert np.array_equal(result.shape, (n_sample,))
-        assert circuit._qfunc_output.shape(dev, Shots(n_sample)) == (n_sample,)
+        assert circuit._qfunc_output.shape(shots=n_sample, num_device_wires=1) == (n_sample,)
 
     def test_multi_wire_sample_regular_shape(self):
         """Test the return type and shape of sampling multiple wires
@@ -96,9 +96,9 @@ class TestSample:
 
         result = circuit()
 
-        assert circuit._qfunc_output[0].shape(dev, Shots(n_sample)) == (n_sample,)
-        assert circuit._qfunc_output[1].shape(dev, Shots(n_sample)) == (n_sample,)
-        assert circuit._qfunc_output[2].shape(dev, Shots(n_sample)) == (n_sample,)
+        assert circuit._qfunc_output[0].shape(shots=n_sample, num_device_wires=3) == (n_sample,)
+        assert circuit._qfunc_output[1].shape(shots=n_sample, num_device_wires=3) == (n_sample,)
+        assert circuit._qfunc_output[2].shape(shots=n_sample, num_device_wires=3) == (n_sample,)
 
         # If all the dimensions are equal the result will end up to be a proper rectangular array
         assert isinstance(result, tuple)
@@ -347,14 +347,12 @@ class TestSample:
 
     def test_shape_no_shots_error(self):
         """Test that the appropriate error is raised with no shots are specified"""
-        dev = qml.device("default.qubit", wires=2, shots=None)
-        shots = Shots(None)
         mp = qml.sample()
 
         with pytest.raises(
             MeasurementShapeError, match="Shots are required to obtain the shape of the measurement"
         ):
-            _ = mp.shape(dev, shots)
+            _ = mp.shape(shots=None)
 
     @pytest.mark.parametrize(
         "obs",
@@ -368,34 +366,15 @@ class TestSample:
     def test_shape(self, obs):
         """Test that the shape is correct."""
         shots = 10
-        dev = qml.device("default.qubit", wires=3, shots=shots)
         res = qml.sample(obs) if obs is not None else qml.sample()
         expected = (shots,) if obs is not None else (shots, 3)
-        assert res.shape(dev, Shots(shots)) == expected
+        assert res.shape(10, 3) == expected
 
     @pytest.mark.parametrize("n_samples", (1, 10))
     def test_shape_wires(self, n_samples):
         """Test that the shape is correct when wires are provided."""
-        dev = qml.device("default.qubit", wires=3, shots=n_samples)
         mp = qml.sample(wires=(0, 1))
-        assert mp.shape(dev, Shots(n_samples)) == (n_samples, 2) if n_samples != 1 else (2,)
-
-    @pytest.mark.parametrize(
-        "obs",
-        [
-            None,
-            qml.PauliZ(0),
-            qml.Hermitian(np.diag([1, 2]), 0),
-            qml.Hermitian(np.diag([1.0, 2.0]), 0),
-        ],
-    )
-    def test_shape_shot_vector(self, obs):
-        """Test that the shape is correct with the shot vector too."""
-        shot_vector = (1, 2, 3)
-        dev = qml.device("default.qubit", wires=3, shots=shot_vector)
-        res = qml.sample(obs) if obs is not None else qml.sample()
-        expected = ((), (2,), (3,)) if obs is not None else ((3,), (2, 3), (3, 3))
-        assert res.shape(dev, Shots(shot_vector)) == expected
+        assert mp.shape(n_samples, 3) == (n_samples, 2) if n_samples != 1 else (2,)
 
     def test_shape_shot_vector_obs(self):
         """Test that the shape is correct with the shot vector and a observable too."""
@@ -494,9 +473,7 @@ def test_jitting_with_sampling_on_subset_of_wires(samples):
 
     expected = (2,) if samples == 1 else (samples, 2)
     assert results.shape == expected
-    assert (
-        circuit._qfunc_output.shape(dev, Shots(samples)) == (samples, 2) if samples != 1 else (2,)
-    )
+    assert circuit._qfunc_output.shape(samples, 3) == (samples, 2) if samples != 1 else (2,)
 
 
 @pytest.mark.jax
diff --git a/tests/measurements/test_state.py b/tests/measurements/test_state.py
index 6dce16479..838cad68e 100644
--- a/tests/measurements/test_state.py
+++ b/tests/measurements/test_state.py
@@ -20,15 +20,7 @@ from pennylane import numpy as pnp
 from pennylane.devices import DefaultMixed
 from pennylane.math.matrix_manipulation import _permute_dense_matrix
 from pennylane.math.quantum import reduce_dm, reduce_statevector
-from pennylane.measurements import (
-    DensityMatrixMP,
-    Shots,
-    State,
-    StateMP,
-    density_matrix,
-    expval,
-    state,
-)
+from pennylane.measurements import DensityMatrixMP, State, StateMP, density_matrix, expval, state
 from pennylane.wires import WireError, Wires
 
 
@@ -521,16 +513,8 @@ class TestState:
     @pytest.mark.parametrize("shots", [None, 1, 10])
     def test_shape(self, shots):
         """Test that the shape is correct for qml.state."""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
         res = qml.state()
-        assert res.shape(dev, Shots(shots)) == (2**3,)
-
-    @pytest.mark.parametrize("s_vec", [(3, 2, 1), (1, 5, 10), (3, 1, 20)])
-    def test_shape_shot_vector(self, s_vec):
-        """Test that the shape is correct for qml.state with the shot vector too."""
-        dev = qml.device("default.qubit", wires=3, shots=s_vec)
-        res = qml.state()
-        assert res.shape(dev, Shots(s_vec)) == ((2**3,), (2**3,), (2**3,))
+        assert res.shape(shots, 3) == (2**3,)
 
     def test_numeric_type(self):
         """Test that the numeric type of state measurements."""
@@ -1091,17 +1075,5 @@ class TestDensityMatrix:
     @pytest.mark.parametrize("shots", [None, 1, 10])
     def test_shape(self, shots):
         """Test that the shape is correct for qml.density_matrix."""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
         res = qml.density_matrix(wires=[0, 1])
-        assert res.shape(dev, Shots(shots)) == (2**2, 2**2)
-
-    @pytest.mark.parametrize("s_vec", [(3, 2, 1), (1, 5, 10), (3, 1, 20)])
-    def test_shape_shot_vector(self, s_vec):
-        """Test that the shape is correct for qml.density_matrix with the shot vector too."""
-        dev = qml.device("default.qubit", wires=3, shots=s_vec)
-        res = qml.density_matrix(wires=[0, 1])
-        assert res.shape(dev, Shots(s_vec)) == (
-            (2**2, 2**2),
-            (2**2, 2**2),
-            (2**2, 2**2),
-        )
+        assert res.shape(shots, 3) == (2**2, 2**2)
diff --git a/tests/measurements/test_var.py b/tests/measurements/test_var.py
index 8d0619a21..e737a6bbe 100644
--- a/tests/measurements/test_var.py
+++ b/tests/measurements/test_var.py
@@ -18,7 +18,7 @@ import pytest
 from flaky import flaky
 
 import pennylane as qml
-from pennylane.measurements import Shots, Variance, VarianceMP
+from pennylane.measurements import Variance, VarianceMP
 
 
 class TestVar:
@@ -156,22 +156,10 @@ class TestVar:
     )
     def test_shape(self, obs):
         """Test that the shape is correct."""
-        dev = qml.device("default.qubit", wires=1)
         res = qml.var(obs)
         # pylint: disable=use-implicit-booleaness-not-comparison
-        assert res.shape(dev, Shots(None)) == ()
-        assert res.shape(dev, Shots(100)) == ()
-
-    @pytest.mark.parametrize(
-        "obs",
-        [qml.PauliZ(0), qml.Hermitian(np.diag([1, 2]), 0), qml.Hermitian(np.diag([1.0, 2.0]), 0)],
-    )
-    def test_shape_shot_vector(self, obs):
-        """Test that the shape is correct with the shot vector too."""
-        res = qml.var(obs)
-        shot_vector = (1, 2, 3)
-        dev = qml.device("default.qubit", wires=3, shots=shot_vector)
-        assert res.shape(dev, Shots(shot_vector)) == ((), (), ())
+        assert res.shape(None, 1) == ()
+        assert res.shape(100, 1) == ()
 
     @pytest.mark.parametrize("state", [np.array([0, 0, 0]), np.array([1, 0, 0, 0, 0, 0, 0, 0])])
     @pytest.mark.parametrize("shots", [None, 1000, [1000, 1111]])
diff --git a/tests/measurements/test_vn_entropy.py b/tests/measurements/test_vn_entropy.py
index 7ab1fbbfc..bd9b199a1 100644
--- a/tests/measurements/test_vn_entropy.py
+++ b/tests/measurements/test_vn_entropy.py
@@ -18,7 +18,7 @@ import numpy as np
 import pytest
 
 import pennylane as qml
-from pennylane.measurements import Shots, VnEntropy
+from pennylane.measurements import VnEntropy
 from pennylane.measurements.vn_entropy import VnEntropyMP
 from pennylane.wires import Wires
 
@@ -109,13 +109,12 @@ class TestInitialization:
         assert meas.numeric_type == float
         assert meas.return_type == VnEntropy
 
-    @pytest.mark.parametrize("shots, shape", [(None, ()), (10, ()), ((1, 10), ((), ()))])
+    @pytest.mark.parametrize("shots, shape", [(None, ()), (10, ())])
     def test_shape(self, shots, shape):
         """Test the ``shape`` method."""
         meas = qml.vn_entropy(wires=0)
-        dev = qml.device("default.qubit", wires=1, shots=shots)
 
-        assert meas.shape(dev, Shots(shots)) == shape
+        assert meas.shape(shots, 1) == shape
 
 
 class TestIntegration:
diff --git a/tests/tape/test_qscript.py b/tests/tape/test_qscript.py
index bcafc1da6..ba9819a37 100644
--- a/tests/tape/test_qscript.py
+++ b/tests/tape/test_qscript.py
@@ -991,83 +991,6 @@ multi_measurements = [
 ]
 
 
-class TestMeasurementProcess:
-    """Tests for the shape and numeric type of a measurement process"""
-
-    measurements_no_shots = [
-        (qml.expval(qml.PauliZ(0)), ()),
-        (qml.var(qml.PauliZ(0)), ()),
-        (qml.probs(wires=[0, 1]), (4,)),
-        (qml.state(), (8,)),
-        (qml.density_matrix(wires=[0, 1]), (4, 4)),
-        (qml.mutual_info(wires0=[0], wires1=[1]), ()),
-        (qml.vn_entropy(wires=[0, 1]), ()),
-    ]
-
-    measurements_finite_shots = [
-        (qml.expval(qml.PauliZ(0)), ()),
-        (qml.var(qml.PauliZ(0)), ()),
-        (qml.probs(wires=[0, 1]), (4,)),
-        (qml.state(), (8,)),
-        (qml.density_matrix(wires=[0, 1]), (4, 4)),
-        (qml.sample(qml.PauliZ(0)), (10,)),
-        (qml.sample(), (10, 3)),
-        (qml.mutual_info(wires0=0, wires1=1), ()),
-        (qml.vn_entropy(wires=[0, 1]), ()),
-    ]
-
-    measurements_shot_vector = [
-        (qml.expval(qml.PauliZ(0)), ((), (), ())),
-        (qml.var(qml.PauliZ(0)), ((), (), ())),
-        (qml.probs(wires=[0, 1]), ((4,), (4,), (4,))),
-        (qml.state(), ((8,), (8,), (8,))),
-        (qml.density_matrix(wires=[0, 1]), ((4, 4), (4, 4), (4, 4))),
-        (qml.sample(qml.PauliZ(0)), ((10,), (20,), (30,))),
-        (qml.sample(), ((10, 3), (20, 3), (30, 3))),
-        (qml.mutual_info(wires0=0, wires1=1), ((), (), ())),
-        (qml.vn_entropy(wires=[0, 1]), ((), (), ())),
-    ]
-
-    @pytest.mark.parametrize("measurement, expected_shape", measurements_no_shots)
-    def test_output_shapes_no_shots(self, measurement, expected_shape):
-        """Test that the output shape of the measurement process is expected
-        when shots=None"""
-        num_wires = 3
-        dev = qml.device("default.qubit", wires=num_wires, shots=None)
-
-        assert measurement.shape(dev, Shots(None)) == expected_shape
-
-    @pytest.mark.parametrize("measurement, expected_shape", measurements_finite_shots)
-    def test_output_shapes_finite_shots(self, measurement, expected_shape):
-        """Test that the output shape of the measurement process is expected
-        when shots is finite"""
-        num_wires = 3
-        num_shots = 10
-        dev = qml.device("default.qubit", wires=num_wires, shots=num_shots)
-
-        assert measurement.shape(dev, Shots(num_shots)) == expected_shape
-
-    @pytest.mark.parametrize("measurement, expected_shape", measurements_shot_vector)
-    def test_output_shapes_shot_vector(self, measurement, expected_shape):
-        """Test that the output shape of the measurement process is expected
-        when shots is a vector"""
-        num_wires = 3
-        shot_vector = [10, 20, 30]
-        dev = qml.device("default.qubit", wires=num_wires, shots=shot_vector)
-
-        assert measurement.shape(dev, Shots(shot_vector)) == expected_shape
-
-    def test_undefined_shape_error(self):
-        """Test that an error is raised for a measurement with an undefined shape"""
-        measurement = qml.counts(wires=[0, 1])
-        dev = qml.device("default.qubit", wires=2)
-        shots = Shots(None)
-        msg = "The shape of the measurement CountsMP is not defined"
-
-        with pytest.raises(qml.QuantumFunctionError, match=msg):
-            measurement.shape(dev, shots)
-
-
 class TestOutputShape:
     """Tests for determining the tape output shape of tapes."""
 
@@ -1128,51 +1051,6 @@ class TestOutputShape:
 
         assert qs.shape(dev) == res_shape
 
-    def test_output_shapes_single_qnode_check_cutoff(self):
-        """Test that the tape output shape is correct when computing
-        probabilities with a dummy device that defines a cutoff value."""
-
-        class CustomDevice(qml.QubitDevice):
-            """A dummy device that has a cutoff value specified and returns
-            analytic probabilities in a fashion similar to the
-            strawberryfields.fock device.
-
-            Note: this device definition is used as PennyLane-SF is not a
-            dependency of PennyLane core and there are no CV device in
-            PennyLane core using a cutoff value.
-            """
-
-            name = "Device with cutoff"
-            short_name = "dummy.device"
-            pennylane_requires = "0.1.0"
-            version = "0.0.1"
-            author = "CV quantum"
-
-            operations = {}
-            observables = {"Identity"}
-
-            def __init__(self, shots=None, wires=None, cutoff=None):
-                super().__init__(wires=wires, shots=shots)
-                self.cutoff = cutoff
-
-            def apply(self, operations, **kwargs):
-                pass
-
-            def analytic_probability(self, wires=None):
-                if wires is None:
-                    wires = self.wires
-                return np.zeros(self.cutoff ** len(wires))
-
-        dev = CustomDevice(wires=2, cutoff=13)
-
-        # If PennyLane-SF is installed, the following can be checked e.g., locally:
-        # dev = qml.device("strawberryfields.fock", wires=2, cutoff_dim=13)
-
-        qs = QuantumScript(measurements=[qml.probs(wires=[0])])
-
-        res_shape = qml.execute([qs], dev, gradient_fn=None)[0]
-        assert qs.shape(dev) == res_shape.shape
-
     @pytest.mark.autograd
     @pytest.mark.parametrize("measurements, expected", multi_measurements)
     @pytest.mark.parametrize("shots", [None, 1, 10])
@@ -1374,16 +1252,10 @@ class TestOutputShape:
         broadcasting along with a device with a shot vector raises an error."""
         dev = qml.device("default.qubit", wires=3, shots=(1, 2, 3))
 
-        with qml.queuing.AnnotatedQueue() as q:
-            qml.RY(np.array([0.1, 0.2]), wires=0)
-            qml.RX(np.array([0.3, 0.4]), wires=0)
-            qml.expval(qml.PauliZ(0))
-
-        tape = qml.tape.QuantumScript.from_queue(q, shots=(1, 2, 3))
-        msg = "Parameter broadcasting when using a shot vector is not supported yet"
+        y = np.array([0.1, 0.2])
+        tape = qml.tape.QuantumScript([qml.RY(y, 0)], [qml.expval(qml.Z(0))], shots=(1, 2, 3))
 
-        with pytest.raises(NotImplementedError, match=msg):
-            tape.shape(dev)
+        assert tape.shape(dev) == ((2,), (2,), (2,))
 
 
 class TestNumericType:
diff --git a/tests/tape/test_tape.py b/tests/tape/test_tape.py
index 227d4025a..dd9e6f99c 100644
--- a/tests/tape/test_tape.py
+++ b/tests/tape/test_tape.py
@@ -1892,54 +1892,6 @@ class TestOutputShape:
         res_shape = res_shape if res_shape != tuple() else ()
         assert tape.shape(dev) == res_shape
 
-    def test_output_shapes_single_qnode_check_cutoff(self):
-        """Test that the tape output shape is correct when computing
-        probabilities with a dummy device that defines a cutoff value."""
-
-        class CustomDevice(qml.QubitDevice):
-            """A dummy device that has a cutoff value specified and returns
-            analytic probabilities in a fashion similar to the
-            strawberryfields.fock device.
-
-            Note: this device definition is used as PennyLane-SF is not a
-            dependency of PennyLane core and there are no CV device in
-            PennyLane core using a cutoff value.
-            """
-
-            name = "Device with cutoff"
-            short_name = "dummy.device"
-            pennylane_requires = "0.1.0"
-            version = "0.0.1"
-            author = "CV quantum"
-
-            operations = {}
-            observables = {"Identity"}
-
-            def __init__(self, shots=None, wires=None, cutoff=None):
-                super().__init__(wires=wires, shots=shots)
-                self.cutoff = cutoff
-
-            def apply(self, operations, **kwargs):
-                pass
-
-            def analytic_probability(self, wires=None):
-                if wires is None:
-                    wires = self.wires
-                return np.zeros(self.cutoff ** len(wires))
-
-        dev = CustomDevice(wires=2, cutoff=13)
-
-        # If PennyLane-SF is installed, the following can be checked e.g., locally:
-        # dev = qml.device("strawberryfields.fock", wires=2, cutoff_dim=13)
-
-        with qml.queuing.AnnotatedQueue() as q:
-            qml.probs(wires=[0])
-
-        tape = qml.tape.QuantumScript.from_queue(q)
-
-        res_shape = qml.execute([tape], dev, gradient_fn=qml.gradients.param_shift_cv)[0]
-        assert tape.shape(dev) == res_shape.shape
-
     @pytest.mark.autograd
     @pytest.mark.parametrize("measurements, expected", multi_measurements)
     @pytest.mark.parametrize("shots", [None, 1, 10])
diff --git a/tests/test_qnode.py b/tests/test_qnode.py
index c40079888..d6ac2c8e0 100644
--- a/tests/test_qnode.py
+++ b/tests/test_qnode.py
@@ -644,7 +644,7 @@ class TestTapeConstruction:
         jitted_qnode1 = jax.jit(qn)
 
         with pytest.raises(
-            qml.QuantumFunctionError, match="Can't JIT a quantum function that returns counts."
+            NotImplementedError, match="The JAX-JIT interface doesn't support qml.counts."
         ):
             jitted_qnode1(0.123)
 
@@ -659,7 +659,7 @@ class TestTapeConstruction:
         jitted_qnode2 = jax.jit(circuit2)
 
         with pytest.raises(
-            qml.QuantumFunctionError, match="Can't JIT a quantum function that returns counts."
+            NotImplementedError, match="The JAX-JIT interface doesn't support qml.counts."
         ):
             jitted_qnode2(0.123)
 
diff --git a/tests/test_qnode_legacy.py b/tests/test_qnode_legacy.py
index 3ab93aece..81130fbff 100644
--- a/tests/test_qnode_legacy.py
+++ b/tests/test_qnode_legacy.py
@@ -716,7 +716,7 @@ class TestTapeConstruction:
         jitted_qnode1 = jax.jit(qn)
 
         with pytest.raises(
-            qml.QuantumFunctionError, match="Can't JIT a quantum function that returns counts."
+            NotImplementedError, match="The JAX-JIT interface doesn't support qml.counts."
         ):
             jitted_qnode1(0.123)
 
@@ -731,7 +731,7 @@ class TestTapeConstruction:
         jitted_qnode2 = jax.jit(circuit2)
 
         with pytest.raises(
-            qml.QuantumFunctionError, match="Can't JIT a quantum function that returns counts."
+            NotImplementedError, match="The JAX-JIT interface doesn't support qml.counts."
         ):
             jitted_qnode2(0.123)