Repository: PennyLaneAI/pennylane
Pull request #3399 (2022-11-21 16:15:15)

CMP #14: Stop using `return_type`
Remove the `return_type` attribute where possible.

**Note:** We need to update all the functions/methods that use `Observable.return_type` to make them use the `MeasurementProcess` class before removing the `return_type` attribute completely. These include:
- `Device.execute`
- `QubitDevice._samples_to_counts`
- `_is_observable` in `circuit_graph.py`+ `CircuitGraph.__init__`
- `_expval_param_shift_tuple` in `parameter_shift.py`

codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3399?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3399](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3399?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (39ad53c) into [remove-return-type-2](https://codecov.io/gh/PennyLaneAI/pennylane/commit/a40a61dd025c18e742640226d9e44e12551f7b44?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (a40a61d) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@                  Coverage Diff                  @@
##           remove-return-type-2    #3399   +/-   ##
=====================================================
  Coverage                 99.70%   99.70%           
=====================================================
  Files                       298      298           
  Lines                     26430    26435    +5     
=====================================================
+ Hits                      26351    26356    +5     
  Misses                       79       79           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3399?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/devices/default\_mixed.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3399/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9taXhlZC5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/fourier/qnode\_spectrum.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3399/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ZvdXJpZXIvcW5vZGVfc3BlY3RydW0ucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/gradients/finite\_difference.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3399/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9maW5pdGVfZGlmZmVyZW5jZS5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/gradients/jvp.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3399/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9qdnAucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/gradients/parameter\_shift.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3399/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9wYXJhbWV0ZXJfc2hpZnQucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/gradients/parameter\_shift\_cv.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3399/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9wYXJhbWV0ZXJfc2hpZnRfY3YucHk=) | `99.59% <100.00%> (+<0.01%)` | :arrow_up: |
| [pennylane/gradients/parameter\_shift\_hessian.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3399/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9wYXJhbWV0ZXJfc2hpZnRfaGVzc2lhbi5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/interfaces/autograd.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3399/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYXV0b2dyYWQucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/interfaces/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3399/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvamF4LnB5) | `100.00% <100.00%> (ø)` | |
| [pennylane/interfaces/jax\_jit.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3399/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvamF4X2ppdC5weQ==) | `100.00% <100.00%> (ø)` | |
| ... and [6 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3399/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)


Patch:
diff --git a/pennylane/_device.py b/pennylane/_device.py
index 92f91ee58..8dda44164 100644
--- a/pennylane/_device.py
+++ b/pennylane/_device.py
@@ -25,22 +25,20 @@ from functools import lru_cache
 import numpy as np
 
 import pennylane as qml
-from pennylane.operation import (
-    Operation,
-    Observable,
-    Tensor,
-)
 from pennylane.measurements import (
+    Expectation,
+    Probability,
     Sample,
     State,
     Variance,
-    Expectation,
-    Probability,
-    MidMeasure,
-    ShadowExpval,
+    _Counts,
+    _MidMeasure,
+    _Probability,
+    _Sample,
+    _ShadowExpval,
 )
-from pennylane.wires import Wires, WireError
-
+from pennylane.operation import Observable, Operation, Tensor
+from pennylane.wires import WireError, Wires
 
 ShotTuple = namedtuple("ShotTuple", ["shots", "copies"])
 """tuple[int, int]: Represents copies of a shot number."""
@@ -736,9 +734,8 @@ class Device(abc.ABC):
 
         hamiltonian_in_obs = "Hamiltonian" in [obs.name for obs in circuit.observables]
 
-        return_types = [m.return_type for m in circuit.observables]
+        is_shadow = any(isinstance(m, _ShadowExpval) for m in circuit.measurements)
 
-        is_shadow = ShadowExpval in return_types
         hamiltonian_unusable = not supports_hamiltonian or (finite_shots and not is_shadow)
 
         if hamiltonian_in_obs and (hamiltonian_unusable or (use_grouping and grouping_known)):
@@ -757,13 +754,7 @@ class Device(abc.ABC):
             len(circuit._obs_sharing_wires) > 0
             and not hamiltonian_in_obs
             and all(
-                t not in return_types
-                for t in [
-                    qml.measurements.Sample,
-                    qml.measurements.Probability,
-                    qml.measurements.Counts,
-                    qml.measurements.AllCounts,
-                ]
+                not isinstance(m, (_Sample, _Probability, _Counts)) for m in circuit.measurements
             )
         ):
             # Check for case of non-commuting terms and that there are no Hamiltonians
@@ -965,7 +956,7 @@ class Device(abc.ABC):
 
             operation_name = o.name
 
-            if getattr(o, "return_type", None) == MidMeasure and not self.capabilities().get(
+            if isinstance(o, _MidMeasure) and not self.capabilities().get(
                 "supports_mid_measure", False
             ):
                 raise DeviceError(
diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index baa52b964..3a6951b47 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -1775,11 +1775,11 @@ class QubitDevice(Device):
 
         outcomes = []
 
-        if isinstance(obs, MeasurementProcess):
+        if isinstance(obs, _Counts):
             # convert samples and outcomes (if using) from arrays to str for dict keys
             samples = ["".join([str(s.item()) for s in sample]) for sample in samples]
 
-            if obs.return_type is AllCounts:
+            if obs.all_outcomes:
                 outcomes = self.generate_basis_states(num_wires)
                 outcomes = ["".join([str(o.item()) for o in outcome]) for outcome in outcomes]
         elif obs.return_type is AllCounts:
diff --git a/pennylane/circuit_graph.py b/pennylane/circuit_graph.py
index 9a522a745..018998088 100644
--- a/pennylane/circuit_graph.py
+++ b/pennylane/circuit_graph.py
@@ -22,7 +22,7 @@ import numpy as np
 import retworkx as rx
 
 import pennylane as qml
-from pennylane.measurements import Sample, State, _State
+from pennylane.measurements import _Sample, _State
 from pennylane.wires import Wires
 
 
@@ -114,20 +114,18 @@ class CircuitGraph:
         self.num_wires = len(wires)
         """int: number of wires the circuit contains"""
         for k, op in enumerate(queue):
-            op.queue_idx = k  # store the queue index in the Operator
             meas_wires = wires or None  # cannot use empty wire list in MeasurementProcess
-            if hasattr(op, "return_type"):
-                if op.return_type is State:
-                    # State measurements contain no wires by default, but wires are
-                    # required for the circuit drawer, so we recreate the state
-                    # measurement with all wires
-                    op = _State(wires=meas_wires)
-
-                elif op.return_type is Sample and op.wires == Wires([]):
-                    # Sampling without specifying wires is treated as sampling all wires
-                    op = qml.sample(wires=meas_wires)
-
-                op.queue_idx = k
+            if isinstance(op, _State):
+                # State measurements contain no wires by default, but wires are
+                # required for the circuit drawer, so we recreate the state
+                # measurement with all wires
+                op = _State(wires=meas_wires)
+
+            elif isinstance(op, _Sample) and op.wires == Wires([]):
+                # Sampling without specifying wires is treated as sampling all wires
+                op = qml.sample(wires=meas_wires)
+
+            op.queue_idx = k  # store the queue index in the Operator
 
             for w in op.wires:
                 # get the index of the wire on the device
diff --git a/pennylane/devices/default_mixed.py b/pennylane/devices/default_mixed.py
index ceb4049c2..5f134c6b3 100644
--- a/pennylane/devices/default_mixed.py
+++ b/pennylane/devices/default_mixed.py
@@ -34,7 +34,7 @@ from pennylane import (
     Snapshot,
 )
 from pennylane import numpy as np
-from pennylane.measurements import Counts, AllCounts, MutualInfo, Sample, State, VnEntropy
+from pennylane.measurements import _Counts, _MutualInfo, _Sample, _State, _VnEntropy
 from pennylane.operation import Channel
 from pennylane.ops.qubit.attributes import diagonal_in_z_basis
 from pennylane.wires import Wires
@@ -581,23 +581,25 @@ class DefaultMixed(QubitDevice):
 
         if self.readout_err:
             wires_list = []
-            for obs in circuit.observables:
-                if obs.return_type is State:
+            for m in circuit.measurements:
+                if isinstance(m, _State):
                     # State: This returns pre-rotated state, so no readout error.
                     # Assumed to only be allowed if it's the only measurement.
                     self.measured_wires = []
                     return super().execute(circuit, **kwargs)
-                if obs.return_type in (Sample, Counts, AllCounts):
-                    if obs.name == "Identity" and obs.wires in (qml.wires.Wires([]), self.wires):
-                        # Sample, Counts: Readout error applied to all device wires when wires
-                        # not specified or all wires specified.
-                        self.measured_wires = self.wires
-                        return super().execute(circuit, **kwargs)
-                if obs.return_type in (VnEntropy, MutualInfo):
+                if isinstance(m, (_Sample, _Counts)) and m.wires in (
+                    qml.wires.Wires([]),
+                    self.wires,
+                ):
+                    # Sample, Counts: Readout error applied to all device wires when wires
+                    # not specified or all wires specified.
+                    self.measured_wires = self.wires
+                    return super().execute(circuit, **kwargs)
+                if isinstance(m, (_VnEntropy, _MutualInfo)):
                     # VnEntropy, MutualInfo: Computed for the state prior to measurement. So, readout
                     # error need not be applied on the corresponding device wires.
                     continue
-                wires_list.append(obs.wires)
+                wires_list.append(m.wires)
             self.measured_wires = qml.wires.Wires.all_wires(wires_list)
         return super().execute(circuit, **kwargs)
 
diff --git a/pennylane/fourier/qnode_spectrum.py b/pennylane/fourier/qnode_spectrum.py
index 4b38745b7..02fea18b2 100644
--- a/pennylane/fourier/qnode_spectrum.py
+++ b/pennylane/fourier/qnode_spectrum.py
@@ -19,6 +19,7 @@ from inspect import signature
 from itertools import product
 
 import numpy as np
+
 import pennylane as qml
 
 from .utils import get_spectrum, join_spectra
@@ -395,9 +396,9 @@ def qnode_spectrum(qnode, encoding_args=None, argnum=None, decimals=8, validatio
         # After construction, check whether invalid operations (for a spectrum)
         # are present in the QNode
         for m in qnode.qtape.measurements:
-            if m.return_type not in {qml.measurements.Expectation, qml.measurements.Probability}:
+            if not isinstance(m, (qml.measurements._Expectation, qml.measurements._Probability)):
                 raise ValueError(
-                    f"The return_type {m.return_type.value} is not supported as it likely does "
+                    f"The measurement {m.__class__.__name__} is not supported as it likely does "
                     "not admit a Fourier spectrum."
                 )
         cjacs = jac_fn(*args, **kwargs)
diff --git a/pennylane/gradients/finite_difference.py b/pennylane/gradients/finite_difference.py
index c6aa54e1f..73e6bdeb9 100644
--- a/pennylane/gradients/finite_difference.py
+++ b/pennylane/gradients/finite_difference.py
@@ -25,14 +25,15 @@ from scipy.special import factorial
 
 import pennylane as qml
 from pennylane._device import _get_num_copies
+from pennylane.measurements import _Probability
 
+from .general_shift_rules import generate_shifted_tapes
 from .gradient_transform import (
-    gradient_transform,
-    grad_method_validation,
     choose_grad_methods,
+    grad_method_validation,
     gradient_analysis,
+    gradient_transform,
 )
-from .general_shift_rules import generate_shifted_tapes
 
 
 @functools.lru_cache(maxsize=None)
@@ -181,7 +182,7 @@ def _all_zero_grad_new(tape, shots=None):
 
     for m in tape.measurements:
         # TODO: Update shape for CV variables
-        if m.return_type is qml.measurements.Probability:
+        if isinstance(m, _Probability):
             shape = 2 ** len(m.wires)
         else:
             shape = ()
@@ -421,7 +422,7 @@ def _finite_diff_new(
         output_dims = []
         # TODO: Update shape for CV variables
         for m in tape.measurements:
-            if m.return_type is qml.measurements.Probability:
+            if isinstance(m, _Probability):
                 output_dims.append(2 ** len(m.wires))
             else:
                 output_dims.append(1)
diff --git a/pennylane/gradients/jvp.py b/pennylane/gradients/jvp.py
index ecca56bf9..0de1d49b4 100644
--- a/pennylane/gradients/jvp.py
+++ b/pennylane/gradients/jvp.py
@@ -16,10 +16,12 @@ This module contains functions for computing the Jacobian vector product
 of tapes.
 """
 from collections.abc import Sequence
+
 import numpy as np
 
 import pennylane as qml
 from pennylane._device import _get_num_copies
+from pennylane.measurements import _Probability
 
 
 def _convert(jac, tangent):
@@ -403,10 +405,7 @@ def batch_jvp(tapes, tangents, gradient_fn, shots=None, reduction="append", grad
 
 def _single_measurement_zero(m, tangent):
     """Aux function to create a zero tensor from a measurement."""
-    if m.return_type is qml.measurements.Probability:
-        dim = 2 ** len(m.wires)
-    else:
-        dim = ()
+    dim = 2 ** len(m.wires) if isinstance(m, _Probability) else ()
     res = qml.math.convert_like(np.zeros(dim), tangent)
     res = qml.math.cast_like(res, tangent)
     return res
diff --git a/pennylane/gradients/parameter_shift.py b/pennylane/gradients/parameter_shift.py
index bbb658b99..60a47e9e4 100644
--- a/pennylane/gradients/parameter_shift.py
+++ b/pennylane/gradients/parameter_shift.py
@@ -24,7 +24,7 @@ import numpy as np
 
 import pennylane as qml
 from pennylane._device import _get_num_copies
-from pennylane.measurements import MutualInfo, State, VnEntropy
+from pennylane.measurements import _MutualInfo, _State, _Variance, _VnEntropy
 
 from .finite_difference import _all_zero_grad_new, _no_trainable_grad_new, finite_diff
 from .general_shift_rules import (
@@ -963,7 +963,7 @@ def _var_param_shift_tuple(
     argnum = argnum or tape.trainable_params
 
     # Determine the locations of any variance measurements in the measurement queue.
-    var_mask = [m.return_type is qml.measurements.Variance for m in tape.measurements]
+    var_mask = [isinstance(m, _Variance) for m in tape.measurements]
     var_indices = np.where(var_mask)[0]
 
     # Get <A>, the expectation value of the tape with unshifted parameters.
@@ -1069,7 +1069,7 @@ def var_param_shift(
     argnum = argnum or tape.trainable_params
 
     # Determine the locations of any variance measurements in the measurement queue.
-    var_mask = [m.return_type is qml.measurements.Variance for m in tape.measurements]
+    var_mask = [isinstance(m, _Variance) for m in tape.measurements]
     var_idx = np.where(var_mask)[0]
 
     # Get <A>, the expectation value of the tape with unshifted parameters.
@@ -1478,7 +1478,7 @@ def _param_shift_new(
         batch_size of the created tapes.
     """
 
-    if any(m.return_type in [State, VnEntropy, MutualInfo] for m in tape.measurements):
+    if any(isinstance(m, (_State, _VnEntropy, _MutualInfo)) for m in tape.measurements):
         raise ValueError(
             "Computing the gradient of circuits that return the state is not supported."
         )
@@ -1525,7 +1525,7 @@ def _param_shift_new(
     if gradient_recipes is None:
         gradient_recipes = [None] * len(argnum)
 
-    if any(m.return_type is qml.measurements.Variance for m in tape.measurements):
+    if any(isinstance(m, _Variance) for m in tape.measurements):
         g_tapes, fn = var_param_shift(tape, argnum, shifts, gradient_recipes, f0, broadcast, shots)
     else:
         g_tapes, fn = expval_param_shift(
@@ -1871,7 +1871,7 @@ def param_shift(
             shots=shots,
         )
 
-    if any(m.return_type in [State, VnEntropy, MutualInfo] for m in tape.measurements):
+    if any(isinstance(m, (_State, _VnEntropy, _MutualInfo)) for m in tape.measurements):
         raise ValueError(
             "Computing the gradient of circuits that return the state is not supported."
         )
@@ -1920,7 +1920,7 @@ def param_shift(
     if gradient_recipes is None:
         gradient_recipes = [None] * len(argnum)
 
-    if any(m.return_type is qml.measurements.Variance for m in tape.measurements):
+    if any(isinstance(m, _Variance) for m in tape.measurements):
         g_tapes, fn = var_param_shift(tape, argnum, shifts, gradient_recipes, f0, broadcast)
     else:
         g_tapes, fn = expval_param_shift(tape, argnum, shifts, gradient_recipes, f0, broadcast)
diff --git a/pennylane/gradients/parameter_shift_cv.py b/pennylane/gradients/parameter_shift_cv.py
index 53ce7e09f..26f65ed98 100644
--- a/pennylane/gradients/parameter_shift_cv.py
+++ b/pennylane/gradients/parameter_shift_cv.py
@@ -23,6 +23,7 @@ from collections.abc import Sequence
 import numpy as np
 
 import pennylane as qml
+from pennylane.measurements import _Expectation, _Probability, _State, _Variance
 
 from .finite_difference import finite_diff
 from .general_shift_rules import generate_shifted_tapes, process_shifts
@@ -60,7 +61,7 @@ def _grad_method(tape, idx):
 
     for m in tape.measurements:
 
-        if (m.return_type is qml.measurements.Probability) or (m.obs.ev_order not in (1, 2)):
+        if isinstance(m, _Probability) or (m.obs.ev_order not in (1, 2)):
             # Higher-order observables (including probability) only support finite differences.
             best.append("F")
             continue
@@ -85,12 +86,12 @@ def _grad_method(tape, idx):
 
         elif m.obs.ev_order == 2:
 
-            if m.return_type is qml.measurements.Expectation:
+            if isinstance(m, _Expectation):
                 # If the observable is second-order, we must use the second-order
                 # CV parameter shift rule
                 best_method = "A2"
 
-            elif m.return_type is qml.measurements.Variance:
+            elif isinstance(m, _Variance):
                 # we only support analytic variance gradients for
                 # first-order observables
                 best_method = "F"
@@ -203,7 +204,7 @@ def var_param_shift(tape, dev_wires, argnum=None, shifts=None, gradient_recipes=
     argnum = argnum or tape.trainable_params
 
     # Determine the locations of any variance measurements in the measurement queue.
-    var_mask = [m.return_type is qml.measurements.Variance for m in tape.measurements]
+    var_mask = [isinstance(m, _Variance) for m in tape.measurements]
     var_idx = np.where(var_mask)[0]
 
     # Get <A>, the expectation value of the tape with unshifted parameters.
@@ -651,7 +652,7 @@ def param_shift_cv(
     """
 
     # perform gradient method validation
-    if any(m.return_type is qml.measurements.State for m in tape.measurements):
+    if any(isinstance(m, _State) for m in tape.measurements):
         raise ValueError(
             "Computing the gradient of circuits that return the state is not supported."
         )
@@ -683,7 +684,7 @@ def param_shift_cv(
         return [], lambda _: np.zeros([tape.output_dim, len(tape.trainable_params)])
 
     method_map = choose_grad_methods(diff_methods, argnum)
-    var_present = any(m.return_type is qml.measurements.Variance for m in tape.measurements)
+    var_present = any(isinstance(m, _Variance) for m in tape.measurements)
 
     unsupported_params = []
     first_order_params = []
diff --git a/pennylane/gradients/parameter_shift_hessian.py b/pennylane/gradients/parameter_shift_hessian.py
index 7f588d95b..19655019c 100644
--- a/pennylane/gradients/parameter_shift_hessian.py
+++ b/pennylane/gradients/parameter_shift_hessian.py
@@ -15,24 +15,22 @@
 This module contains functions for computing the parameter-shift hessian
 of a qubit-based quantum tape.
 """
-import warnings
 import itertools as it
+import warnings
 from collections.abc import Sequence
 
 import pennylane as qml
 from pennylane import numpy as np
+from pennylane.measurements import _Probability, _State, _Variance
 
-from .gradient_transform import (
-    gradient_analysis,
-    grad_method_validation,
-)
-from .hessian_transform import hessian_transform
-from .parameter_shift import _get_operation_recipe
 from .general_shift_rules import (
     _combine_shift_rules,
-    generate_shifted_tapes,
     generate_multishifted_tapes,
+    generate_shifted_tapes,
 )
+from .gradient_transform import grad_method_validation, gradient_analysis
+from .hessian_transform import hessian_transform
+from .parameter_shift import _get_operation_recipe
 
 
 def _process_argnum(argnum, tape):
@@ -193,7 +191,7 @@ def _all_zero_grad_new(tape):
 
     zeros_list = []
     for m in tape.measurements:
-        shape = 2 ** len(m.wires) if m.return_type is qml.measurements.Probability else ()
+        shape = 2 ** len(m.wires) if isinstance(m, _Probability) else ()
 
         zeros = tuple(
             tuple(qml.math.zeros(shape) for _ in range(num_params)) for _ in range(num_params)
@@ -632,14 +630,14 @@ def param_shift_hessian(tape, argnum=None, diagonal_shifts=None, off_diagonal_sh
         )
 
     # Perform input validation before generating tapes.
-    if any(m.return_type is qml.measurements.State for m in tape.measurements):
+    if any(isinstance(m, _State) for m in tape.measurements):
         raise ValueError(
             "Computing the Hessian of circuits that return the state is not supported."
         )
 
     # The parameter-shift Hessian implementation currently doesn't support variance measurements.
     # TODO: Support variances similar to how param_shift does it
-    if any(m.return_type is qml.measurements.Variance for m in tape.measurements):
+    if any(isinstance(m, _Variance) for m in tape.measurements):
         raise ValueError(
             "Computing the Hessian of circuits that return variances is currently not supported."
         )
@@ -847,14 +845,14 @@ def _param_shift_hessian_tuple(
     """
 
     # Perform input validation before generating tapes.
-    if any(m.return_type is qml.measurements.State for m in tape.measurements):
+    if any(isinstance(m, _State) for m in tape.measurements):
         raise ValueError(
             "Computing the Hessian of circuits that return the state is not supported."
         )
 
     # The parameter-shift Hessian implementation currently doesn't support variance measurements.
     # TODO: Support variances similar to how param_shift does it
-    if any(m.return_type is qml.measurements.Variance for m in tape.measurements):
+    if any(isinstance(m, _Variance) for m in tape.measurements):
         raise ValueError(
             "Computing the Hessian of circuits that return variances is currently not supported."
         )
diff --git a/pennylane/interfaces/autograd.py b/pennylane/interfaces/autograd.py
index 11e4c3782..562de7bbf 100644
--- a/pennylane/interfaces/autograd.py
+++ b/pennylane/interfaces/autograd.py
@@ -23,6 +23,7 @@ from autograd.numpy.numpy_boxes import ArrayBox
 
 import pennylane as qml
 from pennylane import numpy as np
+from pennylane.measurements import _Counts
 
 
 def execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_diff=2, mode=None):
@@ -124,10 +125,7 @@ def _execute(
         res, jacs = execute_fn(tapes, **gradient_kwargs)
 
     for i, r in enumerate(res):
-        if any(
-            m.return_type in (qml.measurements.Counts, qml.measurements.AllCounts)
-            for m in tapes[i].measurements
-        ):
+        if any(isinstance(m, _Counts) for m in tapes[i].measurements):
             continue
 
         if isinstance(r, np.ndarray):
diff --git a/pennylane/interfaces/jax.py b/pennylane/interfaces/jax.py
index 535c16dc4..224983281 100644
--- a/pennylane/interfaces/jax.py
+++ b/pennylane/interfaces/jax.py
@@ -21,8 +21,8 @@ import jax
 import jax.numpy as jnp
 
 import pennylane as qml
-from pennylane.measurements import Sample, Probability
 from pennylane.interfaces import InterfaceUnsupportedError
+from pennylane.measurements import _Counts, _Probability, _Sample
 
 dtype = jnp.float64
 
@@ -136,16 +136,16 @@ def _validate_tapes(tapes):
     """
     for t in tapes:
 
-        return_types = [o.return_type for o in t.observables]
-        set_of_return_types = set(return_types)
-        probs_or_sample_measure = Sample in return_types or Probability in return_types
+        measurement_types = [type(m) for m in t.measurements]
+        set_of_return_types = set(measurement_types)
+        probs_or_sample_measure = _Sample in measurement_types or _Probability in measurement_types
         if probs_or_sample_measure and len(set_of_return_types) > 1:
             raise InterfaceUnsupportedError(
                 "Using the JAX interface, sample and probability measurements cannot be mixed with other measurement types."
             )
 
-        if Probability in return_types:
-            set_len_wires = set(len(o.wires) for o in t.observables)
+        if _Probability in measurement_types:
+            set_len_wires = {len(m.wires) for m in t.measurements}
             if len(set_len_wires) > 1:
                 raise InterfaceUnsupportedError(
                     "Using the JAX interface, multiple probability measurements need to have the same number of wires specified."
@@ -174,14 +174,7 @@ def _execute(
         """Auxiliary function to convert the result of a tape to an array,
         unless the tape had Counts measurements that are represented with
         dictionaries. JAX NumPy arrays don't support dictionaries."""
-        return (
-            jnp.array(r)
-            if not any(
-                m.return_type in (qml.measurements.Counts, qml.measurements.AllCounts)
-                for m in tape.measurements
-            )
-            else r
-        )
+        return jnp.array(r) if not any(isinstance(m, _Counts) for m in tape.measurements) else r
 
     @jax.custom_vjp
     def wrapped_exec(params):
@@ -221,8 +214,8 @@ def _execute(
 
             for t in tapes:
                 multi_probs = (
-                    any(o.return_type is Probability for o in t.observables)
-                    and len(t.observables) > 1
+                    any(isinstance(m, _Probability) for m in t.measurements)
+                    and len(t.measurements) > 1
                 )
 
             if multi_probs:
diff --git a/pennylane/interfaces/jax_jit.py b/pennylane/interfaces/jax_jit.py
index 768c574d1..ed01b215f 100644
--- a/pennylane/interfaces/jax_jit.py
+++ b/pennylane/interfaces/jax_jit.py
@@ -19,12 +19,13 @@ to a PennyLane Device class.
 # pylint: disable=too-many-arguments
 import jax
 import jax.numpy as jnp
-
 import numpy as np
 import semantic_version
+
 import pennylane as qml
 from pennylane.interfaces import InterfaceUnsupportedError
 from pennylane.interfaces.jax import _raise_vector_valued_fwd
+from pennylane.measurements import _Probability
 
 dtype = jnp.float64
 
@@ -175,8 +176,8 @@ def _execute(
         if isinstance(gradient_fn, qml.gradients.gradient_transform):
             for t in tapes:
                 multi_probs = (
-                    any(o.return_type is qml.measurements.Probability for o in t.observables)
-                    and len(t.observables) > 1
+                    any(isinstance(m, _Probability) for m in t.measurements)
+                    and len(t.measurements) > 1
                 )
 
                 if multi_probs:
diff --git a/pennylane/interfaces/tensorflow.py b/pennylane/interfaces/tensorflow.py
index 557b6cac4..452ac48d7 100644
--- a/pennylane/interfaces/tensorflow.py
+++ b/pennylane/interfaces/tensorflow.py
@@ -23,6 +23,7 @@ import tensorflow as tf
 from tensorflow.python.eager import context
 
 import pennylane as qml
+from pennylane.measurements import _Counts
 from pennylane._device import _get_num_copies
 
 
@@ -192,10 +193,7 @@ def execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_d
     for i, tape in enumerate(tapes):
         # convert output to TensorFlow tensors
 
-        if any(
-            m.return_type in (qml.measurements.Counts, qml.measurements.AllCounts)
-            for m in tape.measurements
-        ):
+        if any(isinstance(m, _Counts) for m in tape.measurements):
             continue
 
         if isinstance(res[i], np.ndarray):
diff --git a/pennylane/interfaces/tensorflow_autograph.py b/pennylane/interfaces/tensorflow_autograph.py
index 63bd3fbdf..2f962b914 100644
--- a/pennylane/interfaces/tensorflow_autograph.py
+++ b/pennylane/interfaces/tensorflow_autograph.py
@@ -22,15 +22,15 @@ import numpy as np
 import tensorflow as tf
 
 import pennylane as qml
+from pennylane.measurements import _Sample, _State
 from pennylane._device import _get_num_copies
 
-
 from .tensorflow import (
     _compute_vjp,
     _compute_vjp_new,
-    _to_tensors,
-    _res_restructured,
     _jac_restructured,
+    _res_restructured,
+    _to_tensors,
 )
 
 
@@ -103,7 +103,7 @@ def execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_d
 
         if tape.all_sampled:
             output_types.append(tf.int64)
-        elif tape.measurements[0].return_type is qml.measurements.State:
+        elif isinstance(tape.measurements[0], _State):
             output_types.append(tf.complex128)
         else:
             output_types.append(tf.float64)
@@ -320,9 +320,9 @@ def _execute_new(
 
         o_types = []
         for m in tape.measurements:
-            if m.return_type is qml.measurements.Sample:
+            if isinstance(m, _Sample):
                 o_types.append(tf.int64)
-            elif m.return_type is qml.measurements.State:
+            elif isinstance(m, _State):
                 o_types.append(tf.complex128)
             else:
                 o_types.append(tf.float64)
diff --git a/pennylane/interfaces/torch.py b/pennylane/interfaces/torch.py
index ec1cc7136..314b7c243 100644
--- a/pennylane/interfaces/torch.py
+++ b/pennylane/interfaces/torch.py
@@ -20,6 +20,7 @@ import numpy as np
 import torch
 
 import pennylane as qml
+from pennylane.measurements import _Counts
 
 
 def _compute_vjp(dy, jacs, device=None):
@@ -97,10 +98,7 @@ class ExecuteTapes(torch.autograd.Function):
                 # For backwards compatibility, we flatten ragged tape outputs
                 r = np.hstack(r)
 
-            if any(
-                m.return_type in (qml.measurements.Counts, qml.measurements.AllCounts)
-                for m in ctx.tapes[i].measurements
-            ):
+            if any(isinstance(m, _Counts) for m in ctx.tapes[i].measurements):
                 continue
 
             if isinstance(r, (list, tuple)):
diff --git a/pennylane/measurements/__init__.py b/pennylane/measurements/__init__.py
index 3ec7fc903..77b3a45b8 100644
--- a/pennylane/measurements/__init__.py
+++ b/pennylane/measurements/__init__.py
@@ -37,7 +37,7 @@ from .measurements import (
     Variance,
     VnEntropy,
 )
-from .mid_measure import MeasurementValue, MeasurementValueError, measure
+from .mid_measure import MeasurementValue, MeasurementValueError, _MidMeasure, measure
 from .mutual_info import _MutualInfo, mutual_info
 from .probs import _Probability, probs
 from .sample import _Sample, sample
diff --git a/pennylane/measurements/measurements.py b/pennylane/measurements/measurements.py
index 0d69f0a40..0c125ff72 100644
--- a/pennylane/measurements/measurements.py
+++ b/pennylane/measurements/measurements.py
@@ -629,6 +629,20 @@ class MeasurementTransform(MeasurementProcess):
             '''
     """
 
+    method_name = ""
+    """Devices can override the logic of a measurement process by defining a method with the
+    name ``method_name`` of the corresponding class. The method should have the following signature:
+
+    .. code-block:: python
+
+        def method_name(self, qscript: QuantumScript):
+            '''Device's custom measurement implementation.
+
+            Args:
+                qscript: quantum script to transform
+            '''
+    """
+
     @abstractmethod
     def process(self, qscript, device):
         """Process the given quantum script."""
diff --git a/pennylane/ops/functions/equal.py b/pennylane/ops/functions/equal.py
index 5936d78ad..b5df3ab8d 100644
--- a/pennylane/ops/functions/equal.py
+++ b/pennylane/ops/functions/equal.py
@@ -254,7 +254,6 @@ def _equal_measurements(op1: MeasurementProcess, op2: MeasurementProcess, **kwar
 
     """Determine whether two MeasurementProcess objects are equal"""
 
-    return_types_match = op1.return_type == op2.return_type
     if op1.obs is not None and op2.obs is not None:
         observables_match = equal(op1.obs, op2.obs)
     # check obs equality when either one is None (False) or both are None (True)
@@ -263,7 +262,7 @@ def _equal_measurements(op1: MeasurementProcess, op2: MeasurementProcess, **kwar
     wires_match = op1.wires == op2.wires
     eigvals_match = qml.math.allequal(op1.eigvals(), op2.eigvals())
 
-    return return_types_match and observables_match and wires_match and eigvals_match
+    return observables_match and wires_match and eigvals_match
 
 
 @_equal.register
@@ -289,9 +288,8 @@ def _(op1: _MutualInfo, op2: _MutualInfo, **kwargs):
 def _equal_shadow_measurements(op1: _ShadowExpval, op2: _ShadowExpval, **kwargs):
     """Determine whether two ClassicalShadow objects are equal"""
 
-    return_types_match = op1.return_type == op2.return_type
     wires_match = op1.wires == op2.wires
     H_match = op1.H == op2.H
     k_match = op1.k == op2.k
 
-    return return_types_match and wires_match and H_match and k_match
+    return wires_match and H_match and k_match
diff --git a/pennylane/qinfo/transforms.py b/pennylane/qinfo/transforms.py
index 4a5c9eac9..909c23e83 100644
--- a/pennylane/qinfo/transforms.py
+++ b/pennylane/qinfo/transforms.py
@@ -14,10 +14,11 @@
 """QNode transforms for the quantum information quantities."""
 # pylint: disable=import-outside-toplevel, not-callable
 import functools
-import pennylane as qml
 
-from pennylane.transforms import batch_transform, metric_tensor, adjoint_metric_tensor
+import pennylane as qml
 from pennylane.devices import DefaultQubit
+from pennylane.measurements import _State
+from pennylane.transforms import adjoint_metric_tensor, batch_transform, metric_tensor
 
 
 def reduced_dm(qnode, wires):
@@ -53,9 +54,9 @@ def reduced_dm(qnode, wires):
 
     def wrapper(*args, **kwargs):
         qnode.construct(args, kwargs)
-        return_type = qnode.tape.observables[0].return_type
-        if len(qnode.tape.observables) != 1 or not return_type == qml.measurements.State:
-            raise ValueError("The qfunc return type needs to be a state.")
+        measurements = qnode.tape.measurements
+        if len(measurements) != 1 or not isinstance(measurements[0], _State):
+            raise ValueError("The qfunc measurement needs to be State.")
 
         # TODO: optimize given the wires by creating a tape with relevant operations
         state_built = qnode(*args, **kwargs)
@@ -121,9 +122,9 @@ def purity(qnode, wires):
         # Construct tape
         qnode.construct(args, kwargs)
 
-        # Check return type
-        return_type = qnode.tape.observables[0].return_type
-        if len(qnode.tape.observables) != 1 or not return_type == qml.measurements.State:
+        # Check measurement
+        measurements = qnode.tape.measurements
+        if len(measurements) != 1 or not isinstance(measurements[0], _State):
             raise ValueError("The qfunc return type needs to be a state.")
 
         state_built = qnode(*args, **kwargs)
@@ -176,8 +177,8 @@ def vn_entropy(qnode, wires, base=None):
         # If pure state directly return 0.
         if len(wires) == len(qnode.device.wires):
             qnode.construct(args, kwargs)
-            return_type = qnode.tape.observables[0].return_type
-            if len(qnode.tape.observables) != 1 or not return_type == qml.measurements.State:
+            measurements = qnode.tape.measurements
+            if len(measurements) != 1 or not isinstance(measurements[0], _State):
                 raise ValueError("The qfunc return type needs to be a state.")
             density_matrix = qnode(*args, **kwargs)
             if density_matrix.shape == (density_matrix.shape[0],):
diff --git a/pennylane/qnode.py b/pennylane/qnode.py
index 4528dcc4b..abfb8fdb5 100644
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -25,6 +25,7 @@ import autograd
 import pennylane as qml
 from pennylane import Device
 from pennylane.interfaces import INTERFACE_MAP, SUPPORTED_INTERFACES, set_shots
+from pennylane.measurements import ClassicalShadow, _Counts, _MidMeasure
 from pennylane.tape import QuantumScript, make_qscript
 
 
@@ -546,7 +547,7 @@ class QNode:
             )
 
         terminal_measurements = [
-            m for m in self.tape.measurements if m.return_type != qml.measurements.MidMeasure
+            m for m in self.tape.measurements if not isinstance(m, _MidMeasure)
         ]
         if any(ret != m for ret, m in zip(measurement_processes, terminal_measurements)):
             raise qml.QuantumFunctionError(
@@ -576,10 +577,7 @@ class QNode:
         # operations
         # 2. Move this expansion to Device (e.g., default_expand_fn or
         # batch_transform method)
-        if any(
-            getattr(obs, "return_type", None) == qml.measurements.MidMeasure
-            for obs in self.tape.operations
-        ):
+        if any(isinstance(m, _MidMeasure) for m in self.tape.operations):
             self._tape = qml.defer_measurements(self._tape)
 
         if self.expansion_strategy == "device":
@@ -688,10 +686,7 @@ class QNode:
 
             res = res[0]
 
-        if not isinstance(self._qfunc_output, Sequence) and self._qfunc_output.return_type in (
-            qml.measurements.Counts,
-            qml.measurements.AllCounts,
-        ):
+        if not isinstance(self._qfunc_output, Sequence) and isinstance(self._qfunc_output, _Counts):
             if self.device._has_partitioned_shots():
                 return tuple(res)
 
@@ -699,8 +694,7 @@ class QNode:
             return res[0]
 
         if isinstance(self._qfunc_output, Sequence) and any(
-            m.return_type in (qml.measurements.Counts, qml.measurements.AllCounts)
-            for m in self._qfunc_output
+            isinstance(m, _Counts) for m in self._qfunc_output
         ):
 
             # If Counts was returned with other measurements, then apply the
@@ -719,7 +713,7 @@ class QNode:
         ):
             return res
 
-        if self._qfunc_output.return_type is qml.measurements.Shadow:
+        if isinstance(self._qfunc_output, ClassicalShadow):
             # if classical shadows is returned, then don't squeeze the
             # last axis corresponding to the number of qubits
             return qml.math.squeeze(res, axis=0)
diff --git a/pennylane/shadows/transforms.py b/pennylane/shadows/transforms.py
index 8cd93d44f..64eeb7fab 100644
--- a/pennylane/shadows/transforms.py
+++ b/pennylane/shadows/transforms.py
@@ -13,24 +13,26 @@
 # limitations under the License.
 """Classical shadow transforms"""
 
-from itertools import product
-from functools import reduce, wraps
 import warnings
+from functools import reduce, wraps
+from itertools import product
 
 import pennylane as qml
 import pennylane.numpy as np
+from pennylane.measurements import ClassicalShadow
+from pennylane.tape import QuantumScript
 
 
 @qml.batch_transform
-def _replace_obs(tape, obs, *args, **kwargs):
+def _replace_obs(tape: QuantumScript, obs, *args, **kwargs):
     """
     Tape transform to replace the measurement processes with the given one
     """
     # check if the measurement process of the tape is qml.classical_shadow
-    for o in tape.observables:
-        if o.return_type is not qml.measurements.Shadow:
+    for m in tape.measurements:
+        if not isinstance(m, ClassicalShadow):
             raise ValueError(
-                f"Tape measurement must be {qml.measurements.Shadow!r}, got {o.return_type!r}"
+                f"Tape measurement must be ClassicalShadow, got {m.__class__.__name__!r}"
             )
 
     with qml.queuing.AnnotatedQueue() as q:
diff --git a/pennylane/tape/qscript.py b/pennylane/tape/qscript.py
index dcdfab64a..d623db066 100644
--- a/pennylane/tape/qscript.py
+++ b/pennylane/tape/qscript.py
@@ -24,12 +24,15 @@ from typing import List, Union
 
 import pennylane as qml
 from pennylane.measurements import (
-    AllCounts,
-    Counts,
+    ClassicalShadow,
     MeasurementProcess,
-    Sample,
-    Shadow,
-    ShadowExpval,
+    _Counts,
+    _Expectation,
+    _Probability,
+    _Sample,
+    _ShadowExpval,
+    _State,
+    _Variance,
 )
 from pennylane.operation import Observable, Operator
 from pennylane.queuing import AnnotatedQueue, process_queue
@@ -381,7 +384,7 @@ class QuantumScript:
         self.num_wires = len(self.wires)
 
         is_sample_type = [
-            m.return_type in (Sample, Counts, AllCounts, Shadow, ShadowExpval)
+            isinstance(m, (_Sample, _Counts, ClassicalShadow, _ShadowExpval))
             for m in self.measurements
         ]
         self.is_sampled = any(is_sample_type)
@@ -471,11 +474,11 @@ class QuantumScript:
         self._output_dim = 0
         for m in self.measurements:
             # attempt to infer the output dimension
-            if m.return_type is qml.measurements.Probability:
+            if isinstance(m, _Probability):
                 # TODO: what if we had a CV device here? Having the base as
                 # 2 would have to be swapped to the cutoff value
                 self._output_dim += 2 ** len(m.wires)
-            elif m.return_type is not qml.measurements.State:
+            elif not isinstance(m, _State):
                 self._output_dim += 1
         if self.batch_size:
             self._output_dim *= self.batch_size
@@ -710,21 +713,19 @@ class QuantumScript:
         """
         shape = tuple()
 
-        # We know that there's one type of return_type, gather it from the
-        # first one
-        ret_type = mps[0].return_type
-        if ret_type == qml.measurements.State:
+        # We know that there's one type of measurement, gather it from the first one
+        if isinstance(mps[0], _State):
             raise ValueError(
                 "Getting the output shape of a quantum script with multiple state measurements is not supported."
             )
 
         shot_vector = device._shot_vector
         if shot_vector is None:
-            if ret_type in (qml.measurements.Expectation, qml.measurements.Variance):
+            if isinstance(mps[0], (_Expectation, _Variance)):
 
                 shape = (len(mps),)
 
-            elif ret_type == qml.measurements.Probability:
+            elif isinstance(mps[0], _Probability):
 
                 wires_num_set = {len(meas.wires) for meas in mps}
                 same_num_wires = len(wires_num_set) == 1
@@ -741,7 +742,7 @@ class QuantumScript:
                     # measurement processes act on
                     shape = (sum(2 ** len(m.wires) for m in mps),)
 
-            elif ret_type == qml.measurements.Sample:
+            elif isinstance(mps[0], _Sample):
 
                 dim = mps[0].shape(device)
                 shape = (len(mps),) + dim[1:]
@@ -763,15 +764,14 @@ class QuantumScript:
         """
         shape = tuple()
 
-        ret_type = mps[0].return_type
         shot_vector = device._shot_vector
 
         # Shot vector was defined
-        if ret_type in (qml.measurements.Expectation, qml.measurements.Variance):
+        if isinstance(mps[0], (_Expectation, _Variance)):
             num = sum(shottup.copies for shottup in shot_vector)
             shape = (num, len(mps))
 
-        elif ret_type == qml.measurements.Probability:
+        elif isinstance(mps[0], _Probability):
 
             wires_num_set = {len(meas.wires) for meas in mps}
             same_num_wires = len(wires_num_set) == 1
@@ -791,7 +791,7 @@ class QuantumScript:
             dim = mps[0]._get_num_basis_states(len_wires, device)
             shape = sum(s.copies for s in shot_vector), len(mps), dim
 
-        elif ret_type == qml.measurements.Sample:
+        elif isinstance(mps[0], _Sample):
             shape = []
             for shot_val in device.shot_vector:
                 shots = shot_val.shots
@@ -837,7 +837,7 @@ class QuantumScript:
         if len(self.measurements) == 1:
             output_shape = self._single_measurement_shape(self.measurements[0], device)
         else:
-            num_measurements = len({meas.return_type for meas in self.measurements})
+            num_measurements = len({type(meas) for meas in self.measurements})
             if num_measurements == 1:
                 output_shape = self._multi_homogenous_measurement_shape(self.measurements, device)
             else:
@@ -921,7 +921,7 @@ class QuantumScript:
         """
         if qml.active_return():
             return self._numeric_type_new
-        measurement_types = {meas.return_type for meas in self.measurements}
+        measurement_types = {type(meas) for meas in self.measurements}
         if len(measurement_types) > 1:
             raise ValueError(
                 "Getting the numeric type of a quantum script that contains multiple types of measurements is unsupported."
@@ -929,7 +929,7 @@ class QuantumScript:
 
         # Note: if one of the sample measurements contains outputs that
         # are real, then the entire result will be real
-        if list(measurement_types)[0] == qml.measurements.Sample:
+        if measurement_types.pop() is _Sample:
             return next((float for mp in self.measurements if mp.numeric_type is float), int)
 
         return self.measurements[0].numeric_type
diff --git a/pennylane/tape/tape.py b/pennylane/tape/tape.py
index ec7b3b317..e30b9c9ff 100644
--- a/pennylane/tape/tape.py
+++ b/pennylane/tape/tape.py
@@ -19,7 +19,7 @@ import copy
 from threading import RLock
 
 import pennylane as qml
-from pennylane.measurements import AllCounts, Counts, Probability, Sample
+from pennylane.measurements import _Counts, _Probability, _Sample
 from pennylane.operation import DecompositionUndefinedError, Operator
 from pennylane.queuing import AnnotatedQueue, QueuingManager, process_queue
 
@@ -159,8 +159,7 @@ def expand_tape(qscript, depth=1, stop_at=None, expand_measurements=False):
                 qscript._obs_sharing_wires = diag_obs
             except (TypeError, ValueError) as e:
                 if any(
-                    m.return_type in (Probability, Sample, Counts, AllCounts)
-                    for m in qscript.measurements
+                    isinstance(m, (_Probability, _Sample, _Counts)) for m in qscript.measurements
                 ):
                     raise qml.QuantumFunctionError(
                         "Only observables that are qubit-wise commuting "
diff --git a/pennylane/transforms/defer_measurements.py b/pennylane/transforms/defer_measurements.py
index 327767b5a..b20987037 100644
--- a/pennylane/transforms/defer_measurements.py
+++ b/pennylane/transforms/defer_measurements.py
@@ -12,12 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Code for the tape transform implementing the deferred measurement principle."""
-from pennylane.wires import Wires
-
 import pennylane as qml
+from pennylane.measurements import _MidMeasure
 from pennylane.ops.op_math import ctrl
-from pennylane.transforms import qfunc_transform
 from pennylane.queuing import apply
+from pennylane.transforms import qfunc_transform
+from pennylane.wires import Wires
 
 
 @qfunc_transform
@@ -100,10 +100,7 @@ def defer_measurements(tape):
                 f"Cannot apply operations on {op.wires} as the following wires have been measured already: {op_wires_measured}."
             )
 
-        if (
-            isinstance(op, qml.measurements.MeasurementProcess)
-            and op.return_type == qml.measurements.MidMeasure
-        ):
+        if isinstance(op, _MidMeasure):
             measured_wires[op.id] = op.wires[0]
 
         elif op.__class__.__name__ == "Conditional":
diff --git a/pennylane/transforms/hamiltonian_expand.py b/pennylane/transforms/hamiltonian_expand.py
index 213122b8a..03347eedc 100644
--- a/pennylane/transforms/hamiltonian_expand.py
+++ b/pennylane/transforms/hamiltonian_expand.py
@@ -16,6 +16,7 @@ Contains the hamiltonian expand tape transform
 """
 # pylint: disable=protected-access
 import pennylane as qml
+from pennylane.measurements import _Expectation
 
 
 def hamiltonian_expand(tape, group=True):
@@ -121,7 +122,7 @@ def hamiltonian_expand(tape, group=True):
     if (
         len(tape.measurements) != 1
         or not isinstance(hamiltonian := tape.measurements[0].obs, qml.Hamiltonian)
-        or tape.measurements[0].return_type != qml.measurements.Expectation
+        or not isinstance(tape.measurements[0], _Expectation)
     ):
         raise ValueError(
             "Passed tape must end in `qml.expval(H)`, where H is of type `qml.Hamiltonian`"
diff --git a/pennylane/transforms/qcut/qcut.py b/pennylane/transforms/qcut/qcut.py
index 5552ba026..6426a722a 100644
--- a/pennylane/transforms/qcut/qcut.py
+++ b/pennylane/transforms/qcut/qcut.py
@@ -17,6 +17,7 @@ Miscellaneous support for circuit cutting.
 
 
 import uuid
+
 import pennylane as qml
 from pennylane.operation import Operation
 
diff --git a/pennylane/transforms/split_non_commuting.py b/pennylane/transforms/split_non_commuting.py
index 8715f402f..6de412252 100644
--- a/pennylane/transforms/split_non_commuting.py
+++ b/pennylane/transforms/split_non_commuting.py
@@ -18,7 +18,9 @@ Contains the tape transform that splits non-commuting terms
 from functools import reduce
 
 import numpy as np
+
 import pennylane as qml
+from pennylane.measurements import _Probability, _Sample
 
 from .batch_transform import batch_transform
 
@@ -145,17 +147,15 @@ def split_non_commuting(tape):
         dtype=object, requires_grad=True)
 
     """
-    # TODO: allow for samples and probs
-    obs_fn = {qml.measurements.Expectation: qml.expval, qml.measurements.Variance: qml.var}
-
-    obs_list = tape.observables
-    return_types = [m.return_type for m in obs_list]
 
-    if qml.measurements.Sample in return_types or qml.measurements.Probability in return_types:
+    # TODO: allow for samples and probs
+    if any(isinstance(m, (_Sample, _Probability)) for m in tape.measurements):
         raise NotImplementedError(
             "When non-commuting observables are used, only `qml.expval` and `qml.var` are supported."
         )
 
+    obs_list = tape.observables
+
     # If there is more than one group of commuting observables, split tapes
     groups, group_coeffs = qml.pauli.group_observables(obs_list, range(len(obs_list)))
     if len(groups) > 1:
@@ -163,7 +163,9 @@ def split_non_commuting(tape):
         tapes = []
         for group in groups:
             new_tape = tape.__class__(
-                tape._ops, (obs_fn[type](o) for type, o in zip(return_types, group)), tape._prep
+                tape._ops,
+                (m.__class__(obs=o) for m, o in zip(tape.measurements, group)),
+                tape._prep,
             )
 
             tapes.append(new_tape)


Test patch:
diff --git a/tests/fourier/test_qnode_spectrum.py b/tests/fourier/test_qnode_spectrum.py
index c08b269b6..1bfef4fce 100644
--- a/tests/fourier/test_qnode_spectrum.py
+++ b/tests/fourier/test_qnode_spectrum.py
@@ -15,11 +15,14 @@
 Tests for the Fourier spectrum transform.
 """
 from collections import OrderedDict
-import pytest
+
 import numpy as np
+import pytest
+
 import pennylane as qml
 from pennylane import numpy as pnp
-from pennylane.fourier.qnode_spectrum import qnode_spectrum, _process_ids
+from pennylane.fourier.qnode_spectrum import _process_ids, qnode_spectrum
+from pennylane.measurements import _Sample, _State, _Variance
 
 
 def circuit_0(a):
@@ -354,10 +357,10 @@ class TestCircuits:
             qnode_spectrum(circuit)(1.5)
 
     @pytest.mark.parametrize(
-        "measure_fn, measure_args",
-        [(qml.state, ()), (qml.sample, (qml.PauliZ(0),)), (qml.var, (qml.PauliZ(0),))],
+        "measurement",
+        [_State(), _Sample(obs=qml.PauliZ(0)), _Variance(obs=qml.PauliZ(0))],
     )
-    def test_wrong_return_type_error(self, measure_fn, measure_args):
+    def test_wrong_return_type_error(self, measurement):
         """Test that an error is thrown if the QNode has a ``MeasurementProcess``
         with an inadmissable ``return_type``."""
         dev = qml.device("default.qubit", wires=2)
@@ -365,9 +368,9 @@ class TestCircuits:
         @qml.qnode(dev)
         def circuit(x):
             qml.RX(x, wires=0)
-            return [qml.expval(qml.PauliX(1)), measure_fn(*measure_args)]
+            return [qml.expval(qml.PauliX(1)), qml.apply(measurement)]
 
-        with pytest.raises(ValueError, match=f"{measure_fn.__name__} is not supported"):
+        with pytest.raises(ValueError, match=f"{measurement.__class__.__name__} is not supported"):
             qnode_spectrum(circuit)(1.5)
 
 
diff --git a/tests/qinfo/test_reduced_dm.py b/tests/qinfo/test_reduced_dm.py
index e2d55bd46..015d8d0e8 100644
--- a/tests/qinfo/test_reduced_dm.py
+++ b/tests/qinfo/test_reduced_dm.py
@@ -81,13 +81,13 @@ class TestDensityMatrixQNode:
             qml.RZ(0, wires=[0])
             return qml.expval(qml.PauliX(wires=0))
 
-        with pytest.raises(ValueError, match="The qfunc return type needs to be a state"):
+        with pytest.raises(ValueError, match="The qfunc measurement needs to be State"):
             qml.qinfo.reduced_dm(circuit, wires=[0])()
 
     def test_density_matrix_qnode_jax_jit(self, tol):
         """Test reduced_dm jitting for QNode."""
-        from jax import jit
         import jax.numpy as jnp
+        from jax import jit
 
         angle = jnp.array(0.1)
 
diff --git a/tests/returntypes/transforms/test_shadow_transforms_new.py b/tests/returntypes/transforms/test_shadow_transforms_new.py
index 11776a5df..7fa22efd3 100644
--- a/tests/returntypes/transforms/test_shadow_transforms_new.py
+++ b/tests/returntypes/transforms/test_shadow_transforms_new.py
@@ -15,11 +15,12 @@
 """Unit tests for the classical shadows transforms"""
 
 import builtins
+
 import pytest
 
 import pennylane as qml
-from pennylane.measurements.classical_shadow import _ShadowExpval
 from pennylane import numpy as np
+from pennylane.measurements.classical_shadow import _ShadowExpval
 
 
 def hadamard_circuit(wires, shots=10000, interface="autograd"):
@@ -381,6 +382,6 @@ class TestExpvalTransform:
             qml.Hadamard(0)
             return qml.expval(qml.PauliZ(0))
 
-        msg = "Tape measurement must be shadow, got expval"
+        msg = "Tape measurement must be ClassicalShadow, got '_Expectation'"
         with pytest.raises(ValueError, match=msg):
             circuit()
diff --git a/tests/shadow/test_transforms.py b/tests/shadow/test_transforms.py
index 9f97511d7..dd1129b9b 100644
--- a/tests/shadow/test_transforms.py
+++ b/tests/shadow/test_transforms.py
@@ -355,6 +355,6 @@ class TestExpvalTransform:
             qml.Hadamard(0)
             return qml.expval(qml.PauliZ(0))
 
-        msg = "Tape measurement must be shadow, got expval"
+        msg = "Tape measurement must be ClassicalShadow, got '_Expectation'"
         with pytest.raises(ValueError, match=msg):
             circuit()