Repository: PennyLaneAI/pennylane
Pull request #3439 (2022-11-29 11:52:06)

CMP #12: Add support for custom measurements in `QubitDevice`
## `QubitDevice` now supports custom measurements 🎉 

A user/developer can create a custom measurement by defining a class that inherits from `SampleMeasurement`, `StateMeasurement` or `MeasurementTransform` and defining its corresponding method. See an example below:

```python
import pennylane as qml
from pennylane.measurements import SampleMeasurement

class NewMeasurement(SampleMeasurement):
    def process_samples(self, samples, wire_order, shot_range, bin_size):
        return qml.math.sum(samples[..., self.wires])

dev = qml.device("default.qubit", wires=3, shots=1000)

@qml.qnode(dev)
def circuit():
    qml.PauliX(1)
    return NewMeasurement(wires=[0]), NewMeasurement(wires=[1])

print(circuit())
```
```
[   0 1000]
```

codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3439?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3439](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3439?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (dee635e) into [refactor-mp-3](https://codecov.io/gh/PennyLaneAI/pennylane/commit/6a3dcf738a7d36fb46a1c335a66a7141a3182411?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (6a3dcf7) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@              Coverage Diff               @@
##           refactor-mp-3    #3439   +/-   ##
==============================================
  Coverage          99.69%   99.69%           
==============================================
  Files                299      299           
  Lines              26747    26778   +31     
==============================================
+ Hits               26666    26697   +31     
  Misses                81       81           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3439?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/measurements/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3439/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmVtZW50cy9fX2luaXRfXy5weQ==) | `100.00% <ø> (ø)` | |
| [pennylane/\_qubit\_device.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3439/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19xdWJpdF9kZXZpY2UucHk=) | `99.67% <100.00%> (+0.01%)` | :arrow_up: |
| [pennylane/measurements/classical\_shadow.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3439/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% <100.00%> (ø)` | |
| [pennylane/measurements/counts.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3439/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% <100.00%> (ø)` | |
| [pennylane/measurements/expval.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3439/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% <100.00%> (ø)` | |
| [pennylane/measurements/measurements.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3439/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmVtZW50cy9tZWFzdXJlbWVudHMucHk=) | `99.37% <100.00%> (+<0.01%)` | :arrow_up: |
| [pennylane/measurements/mutual\_info.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3439/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% <100.00%> (ø)` | |
| [pennylane/measurements/probs.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3439/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% <100.00%> (ø)` | |
| [pennylane/measurements/sample.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3439/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% <100.00%> (ø)` | |
| [pennylane/measurements/state.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3439/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% <100.00%> (ø)` | |
| ... and [2 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3439/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/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 974c22824..d73f77c51 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -35,7 +35,7 @@
   [(#3408)](https://github.com/PennyLaneAI/pennylane/pull/3408)
 
 * Support custom measurement processes:
-  * `SampleMeasurement`, `StateMeasurement` and `CustomMeasurement` classes have been added.
+  * `SampleMeasurement`, `StateMeasurement` and `MeasurementTransform` classes have been added.
     They contain an abstract method to process samples/quantum state/quantum script.
     [(#3286)](https://github.com/PennyLaneAI/pennylane/pull/3286)
     [(#3388)](https://github.com/PennyLaneAI/pennylane/pull/3388)
@@ -70,6 +70,10 @@
   * Add `_ShadowExpval` class.
     [(#3388)](https://github.com/PennyLaneAI/pennylane/pull/3388)
 
+  * Allow the execution of `SampleMeasurement`, `StateMeasurement` and `MeasurementTransform`
+    measurement processes in `QubitDevice`.
+    [#3439](https://github.com/PennyLaneAI/pennylane/pull/3439)
+
 * 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)
   [(#3378)](https://github.com/PennyLaneAI/pennylane/pull/3378)
@@ -403,7 +407,6 @@
   `qml.qinfo.classical_fisher`, and `qml.qinfo.quantum_fisher` to support the new return types.
   [(#3449)](https://github.com/PennyLaneAI/pennylane/pull/3449)
 
-
 * File `qcut.py` in `qml.transforms` reorganized into multiple files in `qml.transforms.qcut`
   [3413](https://github.com/PennyLaneAI/pennylane/pull/3413)
 
@@ -471,7 +474,7 @@ Deprecations cycles are tracked at [doc/developement/deprecations.rst](https://d
 * `qml.transforms.measurement_grouping` has been deprecated. Use `qml.transforms.hamiltonian_expand` instead.
   [(#3417)](https://github.com/PennyLaneAI/pennylane/pull/3417)
 
-* The ``observables`` argument in ``QubitDevice.statistics`` is deprecated. Please use ``circuit``
+* The `observables` argument in `QubitDevice.statistics` is deprecated. Please use `circuit`
   instead.
   [(#3433)](https://github.com/PennyLaneAI/pennylane/pull/3433)
 
@@ -480,7 +483,6 @@ Deprecations cycles are tracked at [doc/developement/deprecations.rst](https://d
   wanted seed.
   [(#3388)](https://github.com/PennyLaneAI/pennylane/pull/3388)
 
-
 <h3>Documentation</h3>
 
 * Adds developer documentation for the queuing module.
diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index 4ba8cdc98..b8f7d3926 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -23,6 +23,7 @@ import abc
 import contextlib
 import itertools
 import warnings
+from typing import Union
 
 import numpy as np
 
@@ -36,12 +37,15 @@ from pennylane.measurements import (
     Counts,
     Expectation,
     MeasurementProcess,
+    MeasurementTransform,
     MutualInfo,
     Probability,
     Sample,
+    SampleMeasurement,
     Shadow,
     ShadowExpval,
     State,
+    StateMeasurement,
     Variance,
     VnEntropy,
 )
@@ -386,7 +390,7 @@ class QubitDevice(Device):
         if not self.analytic and self._shot_vector is not None:
             results = self._collect_shotvector_results(circuit, counts_exist)
         else:
-            results = self.statistics(circuit)
+            results = self.statistics(circuit=circuit)
 
         if not circuit.is_sampled:
 
@@ -669,8 +673,9 @@ class QubitDevice(Device):
 
         return Wires.all_wires(list_of_wires)
 
+    # pylint: disable=too-many-statements
     def statistics(
-        self, circuit: QuantumScript = None, shot_range=None, bin_size=None, observables=None
+        self, observables=None, shot_range=None, bin_size=None, circuit: QuantumScript = None
     ):
         """Process measurement results from circuit execution and return statistics.
 
@@ -717,20 +722,30 @@ class QubitDevice(Device):
             * Finally, we repeat the measurement statistics for the final 100 shots,
               ``shot_range=[35, 135]``, ``bin_size=100``.
         """
-        if isinstance(circuit, QuantumScript):
-            observables = circuit.observables
-
+        if observables is not None:
+            if isinstance(observables, QuantumScript):
+                measurements = observables.measurements
+            else:
+                warnings.warn(
+                    message="Using a list of observables in ``QubitDevice.statistics`` is "
+                    "deprecated. Please use a ``QuantumScript`` instead.",
+                    category=UserWarning,
+                )
+                measurements = observables
+        elif circuit is not None:
+            measurements = circuit.measurements
         else:
-            warnings.warn(
-                message="The ``observables`` argument in ``QubitDevice.statistics`` is deprecated. "
-                "Please use ``circuit`` instead.",
-                category=UserWarning,
-            )
-            observables = circuit if circuit is not None else observables
+            raise ValueError("Please provide a circuit into the statistics method.")
 
         results = []
 
-        for obs in observables:
+        for m in measurements:
+            # TODO: Remove this when all overriden measurements support the `MeasurementProcess` class
+            if isinstance(m, MeasurementProcess) and m.obs is not None:
+                obs = m.obs
+                obs.return_type = m.return_type
+            else:
+                obs = m
             # Pass instances directly
             if obs.return_type is Expectation:
                 # Appends a result of shape (num_bins,) if bin_size is not None, else a scalar
@@ -759,7 +774,7 @@ class QubitDevice(Device):
                 )
 
             elif obs.return_type is State:
-                if len(observables) > 1:
+                if len(measurements) > 1:
                     raise qml.QuantumFunctionError(
                         "The state or density matrix cannot be returned in combination"
                         " with other return types"
@@ -823,7 +838,7 @@ class QubitDevice(Device):
                 )
 
             elif obs.return_type is Shadow:
-                if len(observables) > 1:
+                if len(measurements) > 1:
                     raise qml.QuantumFunctionError(
                         "Classical shadows cannot be returned in combination"
                         " with other return types"
@@ -831,13 +846,22 @@ class QubitDevice(Device):
                 results.append(self.classical_shadow(obs, circuit))
 
             elif obs.return_type is ShadowExpval:
-                if len(observables) > 1:
+                if len(measurements) > 1:
                     raise qml.QuantumFunctionError(
                         "Classical shadows cannot be returned in combination"
                         " with other return types"
                     )
                 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))
+
+            elif isinstance(m, (SampleMeasurement, StateMeasurement)):
+                results.append(self._measure(m, shot_range=shot_range, bin_size=bin_size))
+
             elif obs.return_type is not None:
                 raise qml.QuantumFunctionError(
                     f"Unsupported return type specified for observable {obs.name}"
@@ -845,6 +869,51 @@ class QubitDevice(Device):
 
         return results
 
+    def _measure(
+        self,
+        measurement: Union[SampleMeasurement, StateMeasurement],
+        shot_range=None,
+        bin_size=None,
+    ):
+        """Compute the corresponding measurement process depending on ``shots`` and the measurement
+        type.
+
+        Args:
+            measurement (Union[SampleMeasurement, StateMeasurement]): measurement process
+            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.
+
+        Raises:
+            ValueError: if the measurement cannot be computed
+
+        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)
+
+            raise ValueError(
+                "Shots must be specified in the device to compute the measurement "
+                f"{measurement.__class__.__name__}"
+            )
+        if isinstance(measurement, StateMeasurement):
+            warnings.warn(
+                f"Requested measurement {measurement.__class__.__name__} with finite shots; the "
+                "returned state information is analytic and is unaffected by sampling. "
+                "To silence this warning, set shots=None on the device.",
+                UserWarning,
+            )
+            return measurement.process_state(state=self.state, wire_order=self.wires)
+        return measurement.process_samples(
+            samples=self._samples, wire_order=self.wires, shot_range=shot_range, bin_size=bin_size
+        )
+
     def _statistics_new(self, circuit: QuantumScript, shot_range=None, bin_size=None):
         """Process measurement results from circuit execution and return statistics.
 
@@ -891,11 +960,16 @@ class QubitDevice(Device):
             * Finally, we repeat the measurement statistics for the final 100 shots,
               ``shot_range=[35, 135]``, ``bin_size=100``.
         """
-        observables = circuit.observables
+        measurements = circuit.measurements
         results = []
 
-        for obs in observables:
-
+        for m in measurements:
+            # TODO: Remove this when all overriden measurements support the `MeasurementProcess` class
+            if m.obs is not None:
+                obs = m.obs
+                obs.return_type = m.return_type
+            else:
+                obs = m
             # 1. Based on the return_type, compute statistics
             # Pass instances directly
             if obs.return_type is Expectation:
@@ -915,7 +989,7 @@ class QubitDevice(Device):
                 result = self.probability(wires=obs.wires, shot_range=shot_range, bin_size=bin_size)
 
             elif obs.return_type is State:
-                if len(observables) > 1:
+                if len(measurements) > 1:
                     raise qml.QuantumFunctionError(
                         "The state or density matrix cannot be returned in combination"
                         " with other return types"
@@ -978,7 +1052,7 @@ class QubitDevice(Device):
                 result = self.mutual_info(wires0=wires0, wires1=wires1, log_base=obs.log_base)
 
             elif obs.return_type is Shadow:
-                if len(observables) > 1:
+                if len(measurements) > 1:
                     raise qml.QuantumFunctionError(
                         "Classical shadows cannot be returned in combination"
                         " with other return types"
@@ -987,7 +1061,7 @@ class QubitDevice(Device):
                 result = self.classical_shadow(obs, circuit=circuit)
 
             elif obs.return_type is ShadowExpval:
-                if len(observables) > 1:
+                if len(measurements) > 1:
                     raise qml.QuantumFunctionError(
                         "Classical shadows cannot be returned in combination"
                         " with other return types"
@@ -995,6 +1069,17 @@ 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)
+
+            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)
             elif obs.return_type is not None:
                 raise qml.QuantumFunctionError(
                     f"Unsupported return type specified for observable {obs.name}"
diff --git a/pennylane/measurements/__init__.py b/pennylane/measurements/__init__.py
index 4c81c51bc..3ec7fc903 100644
--- a/pennylane/measurements/__init__.py
+++ b/pennylane/measurements/__init__.py
@@ -23,6 +23,7 @@ from .measurements import (
     Expectation,
     MeasurementProcess,
     MeasurementShapeError,
+    MeasurementTransform,
     MidMeasure,
     MutualInfo,
     ObservableReturnTypes,
diff --git a/pennylane/measurements/classical_shadow.py b/pennylane/measurements/classical_shadow.py
index 3f489c263..fd4c43c6d 100644
--- a/pennylane/measurements/classical_shadow.py
+++ b/pennylane/measurements/classical_shadow.py
@@ -24,7 +24,7 @@ import numpy as np
 import pennylane as qml
 from pennylane.wires import Wires
 
-from .measurements import CustomMeasurement, MeasurementShapeError, Shadow, ShadowExpval
+from .measurements import MeasurementShapeError, MeasurementTransform, Shadow, ShadowExpval
 
 
 def shadow_expval(H, k=1, seed=None, seed_recipes=True):
@@ -218,7 +218,7 @@ def classical_shadow(wires, seed=None, seed_recipes=True):
     return ClassicalShadow(wires=wires, seed=seed)
 
 
-class ClassicalShadow(CustomMeasurement):
+class ClassicalShadow(MeasurementTransform):
     """Represents a classical shadow measurement process occurring at the end of a
     quantum variational circuit.
 
@@ -231,11 +231,13 @@ class ClassicalShadow(CustomMeasurement):
         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)
 
-    def process(self, tape, device):
+    def process(self, qscript, device):
         """
         Returns the measured bits and recipes in the classical shadow protocol.
 
@@ -255,14 +257,14 @@ class ClassicalShadow(CustomMeasurement):
         Pauli measurements have shape ``(T, n)``.
 
         This implementation is device-agnostic and works by executing single-shot
-        tapes containing randomized Pauli observables. Devices should override this
+        quantum scripts containing randomized Pauli observables. Devices should override this
         if they can offer cleaner or faster implementations.
 
         .. seealso:: :func:`~.classical_shadow`
 
         Args:
-            tape (QuantumScript): the tape to be processed
-            device (Device): the device used to process the tape
+            qscript (QuantumScript): the quantum script to be processed
+            device (Device): the device used to process the quantum script
 
         Returns:
             tensor_like[int]: A tensor with shape ``(2, T, n)``, where the first row represents
@@ -296,7 +298,7 @@ class ClassicalShadow(CustomMeasurement):
                 ]
 
                 device.reset()
-                device.apply(tape.operations, rotations=tape.diagonalizing_gates + rotations)
+                device.apply(qscript.operations, rotations=qscript.diagonalizing_gates + rotations)
 
                 outcomes[t] = device.generate_samples()[0][mapped_wires]
 
@@ -331,7 +333,7 @@ class ClassicalShadow(CustomMeasurement):
         )
 
 
-class _ShadowExpval(CustomMeasurement):
+class _ShadowExpval(MeasurementTransform):
     """Measures the expectation value of an operator using the classical shadow measurement process.
 
     This has the same arguments as the base class MeasurementProcess, plus other additional
@@ -347,25 +349,29 @@ class _ShadowExpval(CustomMeasurement):
         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
         self.k = k
         super().__init__(*args, **kwargs)
 
-    def process(self, tape, device):
+    def process(self, qscript, device):
         """Compute expectation values using classical shadows in a differentiable manner.
 
         Please refer to :func:`~.pennylane.shadow_expval` for detailed documentation.
 
         Args:
-            tape (QuantumScript): the tape to be processed
-            device (Device): the device used to process the tape
+            qscript (QuantumScript): the quantum script to be processed
+            device (Device): the device used to process the quantum script
 
         Returns:
             float: expectation value estimate.
         """
-        bits, recipes = qml.classical_shadow(wires=self.wires, seed=self.seed).process(tape, device)
+        bits, recipes = qml.classical_shadow(wires=self.wires, seed=self.seed).process(
+            qscript, device
+        )
         shadow = qml.shadows.ClassicalShadow(bits, recipes, wire_map=self.wires.tolist())
         return shadow.expval(self.H, self.k)
 
diff --git a/pennylane/measurements/counts.py b/pennylane/measurements/counts.py
index cccd9e2c3..227a270ff 100644
--- a/pennylane/measurements/counts.py
+++ b/pennylane/measurements/counts.py
@@ -147,6 +147,8 @@ def counts(op=None, wires=None, all_outcomes=False):
 class _Counts(SampleMeasurement):
     """Measurement process that returns the samples of a given observable."""
 
+    method_name = "counts"
+
     def __init__(
         self,
         obs: Union[Observable, None] = None,
diff --git a/pennylane/measurements/expval.py b/pennylane/measurements/expval.py
index 724b7c6cb..d63f53623 100644
--- a/pennylane/measurements/expval.py
+++ b/pennylane/measurements/expval.py
@@ -63,6 +63,8 @@ def expval(op: Operator):
 class _Expectation(SampleMeasurement, StateMeasurement):
     """Measurement process that computes the probability of each computational basis state."""
 
+    method_name = "expval"
+
     @property
     def return_type(self):
         return Expectation
diff --git a/pennylane/measurements/measurements.py b/pennylane/measurements/measurements.py
index 973ab6984..a736123c2 100644
--- a/pennylane/measurements/measurements.py
+++ b/pennylane/measurements/measurements.py
@@ -125,6 +125,25 @@ class MeasurementProcess(ABC):
         log_base (float): Base for the logarithm.
     """
 
+    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,
@@ -579,13 +598,28 @@ class StateMeasurement(MeasurementProcess):
         """
 
 
-class CustomMeasurement(MeasurementProcess):
-    """Custom measurement process.
+class MeasurementTransform(MeasurementProcess):
+    """Measurement process that applies a transform into the given quantum script. This transform
+    is carried out inside the gradient black box, thus is not tracked by the gradient transform.
 
     Any class inheriting from this class should define its own ``process`` method, which takes a
-    device instance and a tape and returns the result of the measurement process.
+    device instance and a quantum script and returns the result of the measurement process.
+    """
+
+    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, tape, device):
-        """Process the given tape."""
+    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 b07f60d1f..a3f785c31 100644
--- a/pennylane/measurements/mutual_info.py
+++ b/pennylane/measurements/mutual_info.py
@@ -87,6 +87,8 @@ def mutual_info(wires0, wires1, log_base=None):
 class _MutualInfo(StateMeasurement):
     """Measurement process that returns the mutual information."""
 
+    method_name = "mutual_info"
+
     # pylint: disable=too-many-arguments, unused-argument
     def __init__(
         self,
diff --git a/pennylane/measurements/probs.py b/pennylane/measurements/probs.py
index 6707faafb..a7a79e675 100644
--- a/pennylane/measurements/probs.py
+++ b/pennylane/measurements/probs.py
@@ -117,6 +117,8 @@ def probs(wires=None, op=None):
 class _Probability(SampleMeasurement, StateMeasurement):
     """Measurement process that computes the probability of each computational basis state."""
 
+    method_name = "probability"
+
     @property
     def return_type(self):
         return Probability
diff --git a/pennylane/measurements/sample.py b/pennylane/measurements/sample.py
index 66d6a491c..15c6426fb 100644
--- a/pennylane/measurements/sample.py
+++ b/pennylane/measurements/sample.py
@@ -117,6 +117,8 @@ def sample(op: Union[Observable, None] = None, wires=None):
 class _Sample(SampleMeasurement):
     """Measurement process that returns the samples of a given observable."""
 
+    method_name = "sample"
+
     @property
     def return_type(self):
         return Sample
diff --git a/pennylane/measurements/state.py b/pennylane/measurements/state.py
index e53436738..c8dd9554e 100644
--- a/pennylane/measurements/state.py
+++ b/pennylane/measurements/state.py
@@ -125,6 +125,8 @@ def density_matrix(wires):
 class _State(StateMeasurement):
     """Measurement process that returns the quantum state."""
 
+    method_name = "state"
+
     @property
     def return_type(self):
         return State
diff --git a/pennylane/measurements/var.py b/pennylane/measurements/var.py
index 70a6c539e..9150b6a89 100644
--- a/pennylane/measurements/var.py
+++ b/pennylane/measurements/var.py
@@ -61,6 +61,8 @@ def var(op: Operator):
 class _Variance(SampleMeasurement, StateMeasurement):
     """Measurement process that computes the variance of the supplied observable."""
 
+    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 2fe55c37a..d59afc2b7 100644
--- a/pennylane/measurements/vn_entropy.py
+++ b/pennylane/measurements/vn_entropy.py
@@ -72,6 +72,8 @@ def vn_entropy(wires, log_base=None):
 class _VnEntropy(StateMeasurement):
     """Measurement process that returns the Von Neumann entropy."""
 
+    method_name = "vn_entropy"
+
     # pylint: disable=too-many-arguments, unused-argument
     def __init__(
         self,


Test patch:
diff --git a/tests/measurements/test_measurements.py b/tests/measurements/test_measurements.py
index ae82e3f32..09d521c81 100644
--- a/tests/measurements/test_measurements.py
+++ b/tests/measurements/test_measurements.py
@@ -21,10 +21,13 @@ from pennylane.measurements import (
     Counts,
     Expectation,
     MeasurementProcess,
+    MeasurementTransform,
     MidMeasure,
     Probability,
     Sample,
+    SampleMeasurement,
     State,
+    StateMeasurement,
     Variance,
     _Counts,
     _Expectation,
@@ -430,3 +433,168 @@ class TestDiagonalizingGates:
         expected_classes = [qml.PauliZ, qml.S, qml.Hadamard]
         for op, c in zip(res, expected_classes):
             assert isinstance(op, c)
+
+
+@pytest.mark.parametrize("switch_return", [qml.enable_return, qml.disable_return])
+class TestSampleMeasurement:
+    """Tests for the SampleMeasurement class."""
+
+    def test_custom_sample_measurement(self, switch_return):
+        """Test the execution of a custom sampled measurement."""
+        switch_return()
+
+        class MyMeasurement(SampleMeasurement):
+            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])
+
+        assert qml.math.allequal(circuit(), [1000, 0])
+
+    def test_sample_measurement_without_shots(self, switch_return):
+        """Test that executing a sampled measurement with ``shots=None`` raises an error."""
+        switch_return()
+
+        class MyMeasurement(SampleMeasurement):
+            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)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.PauliX(0)
+            return MyMeasurement(wires=[0]), MyMeasurement(wires=[1])
+
+        with pytest.raises(
+            ValueError, match="Shots must be specified in the device to compute the measurement "
+        ):
+            circuit()
+
+    def test_method_overriden_by_device(self, switch_return):
+        """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])
+
+        circuit.device.test_method = lambda obs, shot_range=None, bin_size=None: 2
+
+        assert qml.math.allequal(circuit(), [2, 2])
+
+
+@pytest.mark.parametrize("switch_return", [qml.enable_return, qml.disable_return])
+class TestStateMeasurement:
+    """Tests for the SampleMeasurement class."""
+
+    def test_custom_state_measurement(self, switch_return):
+        """Test the execution of a custom state measurement."""
+        switch_return()
+
+        class MyMeasurement(StateMeasurement):
+            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()
+
+        assert circuit() == 1
+
+    def test_sample_measurement_with_shots(self, switch_return):
+        """Test that executing a state measurement with shots raises a warning."""
+        switch_return()
+
+        class MyMeasurement(StateMeasurement):
+            def process_state(self, state, wire_order):
+                return qml.math.sum(state)
+
+        dev = qml.device("default.qubit", wires=2, shots=1000)
+
+        @qml.qnode(dev)
+        def circuit():
+            return MyMeasurement()
+
+        with pytest.warns(
+            UserWarning,
+            match="Requested measurement MyMeasurement with finite shots",
+        ):
+            circuit()
+
+    def test_method_overriden_by_device(self, switch_return):
+        """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()
+
+        circuit.device.test_method = lambda obs, shot_range=None, bin_size=None: 2
+
+        assert circuit() == 2
+
+
+@pytest.mark.parametrize("switch_return", [qml.enable_return, qml.disable_return])
+class TestMeasurementTransform:
+    """Tests for the MeasurementTransform class."""
+
+    def test_custom_measurement(self, switch_return):
+        """Test the execution of a custom measurement."""
+        switch_return()
+
+        class MyMeasurement(MeasurementTransform):
+            def process(self, qscript, device):
+                return {device.shots: len(qscript)}
+
+        dev = qml.device("default.qubit", wires=2, shots=1000)
+
+        @qml.qnode(dev)
+        def circuit():
+            return MyMeasurement()
+
+        assert circuit() == {dev.shots: len(circuit.tape)}
+
+    def test_method_overriden_by_device(self, switch_return):
+        """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)
+
+        @qml.qnode(dev)
+        def circuit():
+            return MyMeasurement()
+
+        circuit.device.test_method = lambda qscript: 2
+
+        assert circuit() == 2
diff --git a/tests/test_qubit_device.py b/tests/test_qubit_device.py
index c67f4c1fd..8f23803c7 100644
--- a/tests/test_qubit_device.py
+++ b/tests/test_qubit_device.py
@@ -321,14 +321,24 @@ class TestExtractStatistics:
             dev = mock_qubit_device_extract_stats()
             with pytest.warns(
                 UserWarning,
-                match="The ``observables`` argument in ``QubitDevice.statistics`` is deprecated. ",
+                match="Using a list of observables in ``QubitDevice.statistics`` is",
             ):
                 dev.statistics([])
             with pytest.warns(
                 UserWarning,
-                match="The ``observables`` argument in ``QubitDevice.statistics`` is deprecated. ",
+                match="Using a list of observables in ``QubitDevice.statistics`` is",
             ):
                 dev.statistics(observables=[])
+            qscript = QuantumScript()
+            with pytest.warns(
+                UserWarning,
+                match="Using a list of observables in ``QubitDevice.statistics`` is",
+            ):
+                dev.statistics([], circuit=qscript)
+            with pytest.raises(
+                ValueError, match="Please provide a circuit into the statistics method"
+            ):
+                dev.statistics()
 
     @pytest.mark.parametrize("returntype", [Expectation, Variance, Sample, Probability, State])
     def test_results_created(self, mock_qubit_device_extract_stats, monkeypatch, returntype):