Repository: PennyLaneAI/pennylane
Pull request #3503 (2022-12-09 15:02:45)

Add `measurement_map` attribute to `QubitDevice`
- Remove `method_name` attribute from `MeasurementProcess` (it was added in this release cycle).
- Add `measurement_map` attribute to `QubitDevice`.

Devices can now define the name of the method used to override any measurement process:
```python
import pennylane as qml
from pennylane.devices import DefaultQubit
from pennylane.measurements import SampleMP

class NewDevice(DefaultQubit):
    def __init__(self, shots, wires):
        super().__init__(shots=shots, wires=wires)
        self.measurement_map[SampleMP] = "sample_measurement"

    def sample_measurement(self, measurement, shot_range=None, bin_size=None):
        return 2

dev = NewDevice(wires=2, shots=1000)

@qml.qnode(dev)
def circuit():
    return qml.sample()

print(circuit())
```
Returns
```
2
```

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3503?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3503](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3503?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (d58e03f) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/d9d79c73a9406de86dc76a9bf65f5af46a992331?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (d9d79c7) will **increase** coverage by `0.14%`.
> The diff coverage is `100.00%`.

```diff
@@            Coverage Diff             @@
##           master    #3503      +/-   ##
==========================================
+ Coverage   99.71%   99.86%   +0.14%     
==========================================
  Files         309      309              
  Lines       27452    27417      -35     
==========================================
+ Hits        27375    27381       +6     
+ Misses         77       36      -41     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3503?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/measurements/classical\_shadow.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3503/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmVtZW50cy9jbGFzc2ljYWxfc2hhZG93LnB5) | `100.00% <ø> (ø)` | |
| [pennylane/measurements/counts.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3503/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmVtZW50cy9jb3VudHMucHk=) | `100.00% <ø> (ø)` | |
| [pennylane/measurements/expval.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3503/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmVtZW50cy9leHB2YWwucHk=) | `100.00% <ø> (ø)` | |
| [pennylane/measurements/measurements.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3503/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmVtZW50cy9tZWFzdXJlbWVudHMucHk=) | `100.00% <ø> (ø)` | |
| [pennylane/measurements/mutual\_info.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3503/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmVtZW50cy9tdXR1YWxfaW5mby5weQ==) | `100.00% <ø> (ø)` | |
| [pennylane/measurements/probs.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3503/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmVtZW50cy9wcm9icy5weQ==) | `100.00% <ø> (ø)` | |
| [pennylane/measurements/sample.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3503/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmVtZW50cy9zYW1wbGUucHk=) | `100.00% <ø> (ø)` | |
| [pennylane/measurements/state.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3503/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmVtZW50cy9zdGF0ZS5weQ==) | `100.00% <ø> (ø)` | |
| [pennylane/measurements/var.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3503/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmVtZW50cy92YXIucHk=) | `100.00% <ø> (ø)` | |
| [pennylane/measurements/vn\_entropy.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3503/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmVtZW50cy92bl9lbnRyb3B5LnB5) | `100.00% <ø> (ø)` | |
| ... and [11 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3503/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)

AlbertMitjans: > Hi @AlbertMitjans, looks good! 👍 I like the clean solution here, kudos 😊
> 
> One thing my approval is conditional on is adding a test case for the new return types system - the logic has been added, but I'm not sure if it's tested. Lmk if I missed something though.

Hi @antalszava! Thank you very much for the super quick review!!

In the `tests/measurements/test_measurements.py` file I added a `switch_return` parametrisation which tests for both old/new return_types. I know it's not the usual way of testing it, but since all the tests should return the same results for both old/new return types I didn't want to add an exact copy of this file.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 2510bb036..0df4dd3a6 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -69,7 +69,9 @@
     `VnEntropyMP`, `MutualInfoMP`, `ClassicalShadowMP` and `ShadowExpvalMP` classes.
 
   * Allow the execution of `SampleMeasurement`, `StateMeasurement` and `MeasurementTransform`
-    measurement processes in `QubitDevice`.
+    measurement processes in `QubitDevice`. Also allow devices to override measurement processes by
+    adding a `measurement_map` attribute, which should contain the measurement class as key and
+    the method name as value.
     [(#3286)](https://github.com/PennyLaneAI/pennylane/pull/3286)
     [(#3388)](https://github.com/PennyLaneAI/pennylane/pull/3388)
     [(#3343)](https://github.com/PennyLaneAI/pennylane/pull/3343)
@@ -84,6 +86,7 @@
     [(#3388)](https://github.com/PennyLaneAI/pennylane/pull/3388)
     [(#3439)](https://github.com/PennyLaneAI/pennylane/pull/3439)
     [(#3466)](https://github.com/PennyLaneAI/pennylane/pull/3466)
+    [(#3503)](https://github.com/PennyLaneAI/pennylane/pull/3503)
 
 * Functionality for fetching symbols and geometry of a compound from the PubChem Database using `qchem.mol_data`.
   [(#3289)](https://github.com/PennyLaneAI/pennylane/pull/3289)
@@ -305,7 +308,6 @@
 
 * Improved the performance of executing circuits under the `jax.vmap` transformation, which can now leverage the batch-execution capabilities of some devices. [(#3452)](https://github.com/PennyLaneAI/pennylane/pull/3452)
   
-
 <h4>Return types project</h4>
 
 * The autograd interface for the new return types now supports devices with shot vectors.
@@ -426,6 +428,7 @@
          [-0.383     , -0.1915    ,  0.        ,  0.        ,  0.1915    ],
          [-0.38466667, -0.19233333,  0.        ,  0.        ,  0.19233333]])>
   ```
+
 * Thy PyTorch interface supports the new return system and users can use jacobian and hessian using custom differentiation
   methods (e.g., parameter-shift, finite difference or adjoint).
   [(#3416)](https://github.com/PennyLaneAI/pennylane/pull/3414)
@@ -440,6 +443,7 @@
       qml.CNOT(wires=[0, 1])>
       return qml.expval(qml.PauliZ(0)), qml.probs([0, 1])
   ```
+
   ```pycon
   >>> a = torch.tensor(0.1, requires_grad=True)
   >>> b = torch.tensor(0.2, requires_grad=True)
@@ -479,6 +483,7 @@
   (DeviceArray(5.55111512e-17, dtype=float64, weak_type=True),
   DeviceArray(0., dtype=float64, weak_type=True)))
   ```
+
 * Updated `qml.transforms.split_non_commuting` to support the new return types.
   [(#3414)](https://github.com/PennyLaneAI/pennylane/pull/3414)
 
@@ -496,7 +501,6 @@
   support the new return types.
   [(#3346)](https://github.com/PennyLaneAI/pennylane/pull/3346)
 
-
 <h3>Breaking changes</h3>
 
 * The `log_base` attribute has been moved from `MeasurementProcess` to the new `VnEntropyMP` and
diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index 8c670a336..060e34879 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -23,6 +23,7 @@ import abc
 import contextlib
 import itertools
 import warnings
+from collections import defaultdict
 from typing import Union
 
 import numpy as np
@@ -214,6 +215,49 @@ class QubitDevice(Device):
         "Prod",
     }
 
+    measurement_map = defaultdict(lambda: "")  # e.g. {SampleMP: "sample"}
+    """Mapping used to override the logic of measurement processes. The dictionary maps a
+    measurement class to a string containing the name of a device's method that overrides the
+    measurement process. The method defined by the device should have the following arguments:
+
+    * measurement (MeasurementProcess): measurement to override
+    * shot_range (tuple[int]): 2-tuple of integers specifying the range of samples
+        to use. If not specified, all samples are used.
+    * bin_size (int): Divides the shot range into bins of size ``bin_size``, and
+        returns the measurement statistic separately over each bin. If not
+        provided, the entire shot range is treated as a single bin.
+
+    .. note::
+
+        When overriding the logic of a :class:`~pennylane.measurements.MeasurementTransform`, the
+        method defined by the device should only have a single argument:
+
+        * qscript: quantum script to transform
+
+    **Example:**
+
+    Let's create device that inherits from :class:`~pennylane.devices.DefaultQubit` and overrides the
+    logic of the `qml.sample` measurement. To do so we will need to update the ``measurement_map``
+    dictionary:
+
+    .. code-block:: python
+
+        class NewDevice(DefaultQubit):
+            def __init__(self, wires, shots):
+                super().__init__(wires=wires, shots=shots)
+                self.measurement_map[SampleMP] = "sample_measurement"
+
+            def sample_measurement(self, measurement, shot_range=None, bin_size=None):
+                return 2
+
+    >>> dev = NewDevice(wires=2, shots=1000)
+    >>> @qml.qnode(dev)
+    ... def circuit():
+    ...     return qml.sample()
+    >>> circuit()
+    tensor(2, requires_grad=True)
+    """
+
     def __init__(
         self, wires=1, shots=None, *, r_dtype=np.float64, c_dtype=np.complex128, analytic=None
     ):
@@ -749,9 +793,15 @@ class QubitDevice(Device):
                 obs.return_type = m.return_type
             else:
                 obs = m
+            # Check if there is an overriden version of the measurement process
+            if method := getattr(self, self.measurement_map[type(m)], False):
+                if isinstance(m, MeasurementTransform):
+                    results.append(method(qscript=circuit))
+                else:
+                    results.append(method(m, shot_range=shot_range, bin_size=bin_size))
             # TODO: Remove return_type when `observables` argument is removed from this method
             # Pass instances directly
-            if obs.return_type is Expectation:
+            elif obs.return_type is Expectation:
                 # Appends a result of shape (num_bins,) if bin_size is not None, else a scalar
                 results.append(self.expval(obs, shot_range=shot_range, bin_size=bin_size))
 
@@ -858,10 +908,7 @@ class QubitDevice(Device):
                 results.append(self.shadow_expval(obs, circuit=circuit))
 
             elif isinstance(m, MeasurementTransform):
-                if method := getattr(self, m.method_name, False):
-                    results.append(method(qscript=circuit))
-                else:
-                    results.append(m.process(qscript=circuit, device=self))
+                results.append(m.process(qscript=circuit, device=self))
 
             elif isinstance(m, (SampleMeasurement, StateMeasurement)):
                 results.append(self._measure(m, shot_range=shot_range, bin_size=bin_size))
@@ -896,8 +943,6 @@ class QubitDevice(Device):
         Returns:
             Union[float, dict, list[float]]: result of the measurement
         """
-        if method := getattr(self, measurement.method_name, False):
-            return method(measurement, shot_range=shot_range, bin_size=bin_size)
         if self.shots is None:
             if isinstance(measurement, StateMeasurement):
                 return measurement.process_state(state=self.state, wire_order=self.wires)
@@ -974,9 +1019,15 @@ class QubitDevice(Device):
                 obs.return_type = m.return_type
             else:
                 obs = m
+            # Check if there is an overriden version of the measurement process
+            if method := getattr(self, self.measurement_map[type(m)], False):
+                if isinstance(m, MeasurementTransform):
+                    result = method(qscript=circuit)
+                else:
+                    result = method(m, shot_range=shot_range, bin_size=bin_size)
             # 1. Based on the measurement type, compute statistics
             # Pass instances directly
-            if isinstance(m, ExpectationMP):
+            elif isinstance(m, ExpectationMP):
                 result = self.expval(obs, shot_range=shot_range, bin_size=bin_size)
 
             elif isinstance(m, VarianceMP):
@@ -1072,16 +1123,10 @@ class QubitDevice(Device):
                 result = self.shadow_expval(obs, circuit=circuit)
 
             elif isinstance(m, MeasurementTransform):
-                if method := getattr(self, m.method_name, False):
-                    result = method(qscript=circuit)
-                else:
-                    result = m.process(qscript=circuit, device=self)
+                result = m.process(qscript=circuit, device=self)
 
             elif isinstance(m, (SampleMeasurement, StateMeasurement)):
-                if method := getattr(self, m.method_name, False):
-                    result = method(obs, shot_range=shot_range, bin_size=bin_size)
-                else:
-                    result = self._measure(m, shot_range=shot_range, bin_size=bin_size)
+                result = self._measure(m, shot_range=shot_range, bin_size=bin_size)
 
             elif obs.return_type is not None:
                 raise qml.QuantumFunctionError(
diff --git a/pennylane/measurements/classical_shadow.py b/pennylane/measurements/classical_shadow.py
index b7a49d5b7..18ea5ae99 100644
--- a/pennylane/measurements/classical_shadow.py
+++ b/pennylane/measurements/classical_shadow.py
@@ -233,8 +233,6 @@ class ClassicalShadowMP(MeasurementTransform):
         kwargs (dict[Any, Any]): Additional keyword arguments passed to :class:`~.pennylane.measurements.MeasurementProcess`
     """
 
-    method_name = "classical_shadow"
-
     def __init__(self, *args, seed=None, **kwargs):
         self.seed = seed
         super().__init__(*args, **kwargs)
@@ -350,8 +348,6 @@ class ShadowExpvalMP(MeasurementTransform):
         kwargs (dict[Any, Any]): Additional keyword arguments passed to :class:`~.pennylane.measurements.MeasurementProcess`
     """
 
-    method_name = "shadow_expval"
-
     def __init__(self, *args, H, seed=None, k=1, **kwargs):
         self.seed = seed
         self.H = H
diff --git a/pennylane/measurements/counts.py b/pennylane/measurements/counts.py
index e90b23ded..816c9329f 100644
--- a/pennylane/measurements/counts.py
+++ b/pennylane/measurements/counts.py
@@ -163,8 +163,6 @@ class CountsMP(SampleMeasurement):
             outcomes (default), or whether it will display all possible outcomes for the system
     """
 
-    method_name = "counts"
-
     # pylint: disable=too-many-arguments
     def __init__(
         self,
diff --git a/pennylane/measurements/expval.py b/pennylane/measurements/expval.py
index 2eebeef24..27c26ef38 100644
--- a/pennylane/measurements/expval.py
+++ b/pennylane/measurements/expval.py
@@ -75,8 +75,6 @@ class ExpectationMP(SampleMeasurement, StateMeasurement):
             where the instance has to be identified
     """
 
-    method_name = "expval"
-
     @property
     def return_type(self):
         return Expectation
diff --git a/pennylane/measurements/measurements.py b/pennylane/measurements/measurements.py
index ff4cc3b87..be129ea2f 100644
--- a/pennylane/measurements/measurements.py
+++ b/pennylane/measurements/measurements.py
@@ -122,25 +122,6 @@ class MeasurementProcess(ABC):
             where the instance has to be identified
     """
 
-    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, measurement: MeasurementProcess, shot_range=None, bin_size=None):
-            '''Device's custom measurement implementation.
-
-            Args:
-                measurement (MeasurementProcess): measurement to override
-                shot_range (tuple[int]): 2-tuple of integers specifying the range of samples
-                    to use. If not specified, all samples are used.
-                bin_size (int): Divides the shot range into bins of size ``bin_size``, and
-                    returns the measurement statistic separately over each bin. If not
-                    provided, the entire shot range is treated as a single bin.
-            '''
-    """
-
     # pylint: disable=too-many-arguments
     def __init__(
         self,
@@ -662,20 +643,6 @@ class MeasurementTransform(MeasurementProcess):
     * device (Device): device used to transform the quantum script
     """
 
-    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/measurements/mutual_info.py b/pennylane/measurements/mutual_info.py
index c5dd937be..ef7492a68 100644
--- a/pennylane/measurements/mutual_info.py
+++ b/pennylane/measurements/mutual_info.py
@@ -106,8 +106,6 @@ class MutualInfoMP(StateMeasurement):
 
     """
 
-    method_name = "mutual_info"
-
     # pylint: disable=too-many-arguments
     def __init__(
         self,
diff --git a/pennylane/measurements/probs.py b/pennylane/measurements/probs.py
index 577452a12..b245c83a0 100644
--- a/pennylane/measurements/probs.py
+++ b/pennylane/measurements/probs.py
@@ -131,8 +131,6 @@ class ProbabilityMP(SampleMeasurement, StateMeasurement):
             where the instance has to be identified
     """
 
-    method_name = "probability"
-
     @property
     def return_type(self):
         return Probability
diff --git a/pennylane/measurements/sample.py b/pennylane/measurements/sample.py
index c153097ae..254d13e8e 100644
--- a/pennylane/measurements/sample.py
+++ b/pennylane/measurements/sample.py
@@ -132,8 +132,6 @@ class SampleMP(SampleMeasurement):
             where the instance has to be identified
     """
 
-    method_name = "sample"
-
     @property
     def return_type(self):
         return Sample
diff --git a/pennylane/measurements/state.py b/pennylane/measurements/state.py
index 4778a867d..fb93b280e 100644
--- a/pennylane/measurements/state.py
+++ b/pennylane/measurements/state.py
@@ -142,8 +142,6 @@ class StateMP(StateMeasurement):
             where the instance has to be identified
     """
 
-    method_name = "state"
-
     @property
     def return_type(self):
         return State
diff --git a/pennylane/measurements/var.py b/pennylane/measurements/var.py
index a51b782df..64a70231b 100644
--- a/pennylane/measurements/var.py
+++ b/pennylane/measurements/var.py
@@ -75,8 +75,6 @@ class VarianceMP(SampleMeasurement, StateMeasurement):
             where the instance has to be identified
     """
 
-    method_name = "var"
-
     @property
     def return_type(self):
         return Variance
diff --git a/pennylane/measurements/vn_entropy.py b/pennylane/measurements/vn_entropy.py
index fe4578056..3e9a87ba4 100644
--- a/pennylane/measurements/vn_entropy.py
+++ b/pennylane/measurements/vn_entropy.py
@@ -89,8 +89,6 @@ class VnEntropyMP(StateMeasurement):
             where the instance has to be identified
     """
 
-    method_name = "vn_entropy"
-
     # pylint: disable=too-many-arguments, unused-argument
     def __init__(
         self,


Test patch:
diff --git a/pennylane/devices/tests/test_measurements.py b/pennylane/devices/tests/test_measurements.py
index c4efcfd11..726e22545 100755
--- a/pennylane/devices/tests/test_measurements.py
+++ b/pennylane/devices/tests/test_measurements.py
@@ -19,7 +19,14 @@ from scipy.sparse import csr_matrix
 
 import pennylane as qml
 from pennylane import numpy as np
-from pennylane.measurements import MeasurementTransform, SampleMeasurement, StateMeasurement
+from pennylane.measurements import (
+    ClassicalShadowMP,
+    MeasurementTransform,
+    SampleMeasurement,
+    SampleMP,
+    StateMeasurement,
+    StateMP,
+)
 
 pytestmark = pytest.mark.skip_unsupported
 
@@ -1530,19 +1537,18 @@ class TestSampleMeasurement:
         """Test that the device can override a measurement process."""
         dev = device(2)
 
-        class MyMeasurement(SampleMeasurement):
-            """Dummy sampled measurement."""
-
-            method_name = "test_method"
-
-            def process_samples(self, samples, wire_order, shot_range=None, bin_size=None):
-                return qml.math.sum(samples[..., self.wires])
+        if dev.shots is None:
+            pytest.skip(
+                "The number of shots has to be explicitly set on the device when using "
+                "sample-based measurements."
+            )
 
         @qml.qnode(dev)
         def circuit():
             qml.PauliX(0)
-            return MyMeasurement(wires=[0]), MyMeasurement(wires=[1])
+            return qml.sample(wires=0), qml.sample(wires=1)
 
+        circuit.device.measurement_map[SampleMP] = "test_method"
         circuit.device.test_method = lambda obs, shot_range=None, bin_size=None: 2
 
         assert qml.math.allequal(circuit(), [2, 2])
@@ -1594,18 +1600,12 @@ class TestStateMeasurement:
         """Test that the device can override a measurement process."""
         dev = device(2)
 
-        class MyMeasurement(StateMeasurement):
-            """Dummy state measurement."""
-
-            method_name = "test_method"
-
-            def process_state(self, state, wire_order):
-                return qml.math.sum(state)
-
         @qml.qnode(dev)
         def circuit():
-            return MyMeasurement()
+            qml.PauliX(0)
+            return qml.state()
 
+        circuit.device.measurement_map[StateMP] = "test_method"
         circuit.device.test_method = lambda obs, shot_range=None, bin_size=None: 2
 
         assert circuit() == 2
@@ -1634,18 +1634,17 @@ class TestCustomMeasurement:
         """Test that the device can override a measurement process."""
         dev = device(2)
 
-        class MyMeasurement(MeasurementTransform):
-            """Dummy measurement transform."""
-
-            method_name = "test_method"
-
-            def process(self, qscript, device):
-                return {device.shots: len(qscript)}
+        if dev.shots is None:
+            pytest.skip(
+                "The number of shots has to be explicitly set on the device when using "
+                "sample-based measurements."
+            )
 
         @qml.qnode(dev)
         def circuit():
-            return MyMeasurement()
+            return qml.classical_shadow(wires=0)
 
+        circuit.device.measurement_map[ClassicalShadowMP] = "test_method"
         circuit.device.test_method = lambda qscript: 2
 
         assert circuit() == 2
diff --git a/tests/measurements/test_measurements.py b/tests/measurements/test_measurements.py
index 7606e21d2..b1464e106 100644
--- a/tests/measurements/test_measurements.py
+++ b/tests/measurements/test_measurements.py
@@ -486,19 +486,14 @@ class TestSampleMeasurement:
         """Test that the device can override a measurement process."""
         switch_return()
 
-        class MyMeasurement(SampleMeasurement):
-            method_name = "test_method"
-
-            def process_samples(self, samples, wire_order, shot_range, bin_size):
-                return qml.math.sum(samples[..., self.wires])
-
         dev = qml.device("default.qubit", wires=2, shots=1000)
 
         @qml.qnode(dev)
         def circuit():
             qml.PauliX(0)
-            return MyMeasurement(wires=[0]), MyMeasurement(wires=[1])
+            return qml.sample(wires=[0]), qml.sample(wires=[1])
 
+        circuit.device.measurement_map[SampleMP] = "test_method"
         circuit.device.test_method = lambda obs, shot_range=None, bin_size=None: 2
 
         assert qml.math.allequal(circuit(), [2, 2])
@@ -548,18 +543,13 @@ class TestStateMeasurement:
         """Test that the device can override a measurement process."""
         switch_return()
 
-        class MyMeasurement(StateMeasurement):
-            method_name = "test_method"
-
-            def process_state(self, state, wire_order):
-                return qml.math.sum(state)
-
         dev = qml.device("default.qubit", wires=2)
 
         @qml.qnode(dev)
         def circuit():
-            return MyMeasurement()
+            return qml.state()
 
+        circuit.device.measurement_map[StateMP] = "test_method"
         circuit.device.test_method = lambda obs, shot_range=None, bin_size=None: 2
 
         assert circuit() == 2
@@ -589,18 +579,13 @@ class TestMeasurementTransform:
         """Test that the device can override a measurement process."""
         switch_return()
 
-        class MyMeasurement(MeasurementTransform):
-            method_name = "test_method"
-
-            def process(self, qscript, device):
-                return {device.shots: len(qscript)}
-
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit", wires=2, shots=1000)
 
         @qml.qnode(dev)
         def circuit():
-            return MyMeasurement()
+            return qml.classical_shadow(wires=0)
 
+        circuit.device.measurement_map[ClassicalShadowMP] = "test_method"
         circuit.device.test_method = lambda qscript: 2
 
         assert circuit() == 2