Repository: PennyLaneAI/pennylane
Pull request #3292 (2022-11-08 15:06:20)

CMP #6: Add `Counts` MP class
1. Create `_Counts` class (name will be changed to `Counts` after the whole refactor).
2. Return `_Counts` instance in `qml.counts`.
3. Copy code from `QubitDevice` into `_Counts.process` + improve code style.
4. Add `teardown_method` to the `test_counts.py` tests that checks that `QubitDevice.sample(counts=True)` and `_Counts.process` return the same values for all tests.

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/3292?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3292](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3292?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (7127ec3) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/1e5083d1e549791cf4be140013701c1e026644eb?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (1e5083d) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master    #3292   +/-   ##
=======================================
  Coverage   99.70%   99.70%           
=======================================
  Files         299      299           
  Lines       26114    26178   +64     
=======================================
+ Hits        26037    26101   +64     
  Misses         77       77           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3292?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/\_qubit\_device.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3292/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19xdWJpdF9kZXZpY2UucHk=) | `99.66% <100.00%> (+<0.01%)` | :arrow_up: |
| [pennylane/circuit\_graph.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3292/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2NpcmN1aXRfZ3JhcGgucHk=) | `93.06% <100.00%> (ø)` | |
| [pennylane/measurements/classical\_shadow.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3292/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/3292/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/measurements.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3292/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% <100.00%> (ø)` | |
| [pennylane/measurements/sample.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3292/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/ops/functions/equal.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3292/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9mdW5jdGlvbnMvZXF1YWwucHk=) | `100.00% <100.00%> (ø)` | |

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 ad5e62c72..3e82c2de8 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -12,6 +12,9 @@
   * Add `_Sample` class.
     [#3288](https://github.com/PennyLaneAI/pennylane/pull/3288)
 
+  * Add `_Counts` class.
+    [#3292](https://github.com/PennyLaneAI/pennylane/pull/3292)
+
   * Add `_State` class.
     [#3287](https://github.com/PennyLaneAI/pennylane/pull/3287)
 
@@ -24,7 +27,7 @@
 * 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)
- 
+
   ```pycon
   >>> mol_data("BeH2")
   (['Be', 'H', 'H'],
diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index a85b69dfe..58ec33fff 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -1074,7 +1074,8 @@ class QubitDevice(Device):
             )
         return np.random.choice(basis_states, shots, p=state_probability)
 
-    def generate_basis_states(self, num_wires, dtype=np.uint32):
+    @staticmethod
+    def generate_basis_states(num_wires, dtype=np.uint32):
         """
         Generates basis states in binary representation according to the number
         of wires specified.
@@ -1102,7 +1103,7 @@ class QubitDevice(Device):
         """
         if 2 < num_wires < 32:
             states_base_ten = np.arange(2**num_wires, dtype=dtype)
-            return self.states_to_binary(states_base_ten, num_wires, dtype=dtype)
+            return QubitDevice.states_to_binary(states_base_ten, num_wires, dtype=dtype)
 
         # A slower, but less memory intensive method
         basis_states_generator = itertools.product((0, 1), repeat=num_wires)
diff --git a/pennylane/measurements/counts.py b/pennylane/measurements/counts.py
index 21f7174f1..b423aa4c0 100644
--- a/pennylane/measurements/counts.py
+++ b/pennylane/measurements/counts.py
@@ -15,11 +15,16 @@
 """
 This module contains the qml.counts measurement.
 """
+# pylint: disable=too-many-arguments, abstract-method
+import copy
 import warnings
+from typing import Sequence, Tuple, Union
 
+import pennylane as qml
+from pennylane.operation import Observable
 from pennylane.wires import Wires
 
-from .measurements import AllCounts, Counts, MeasurementProcess
+from .measurements import AllCounts, Counts, ObservableReturnTypes, SampleMeasurement
 
 
 def counts(op=None, wires=None, all_outcomes=False):
@@ -135,7 +140,119 @@ def counts(op=None, wires=None, all_outcomes=False):
             )
         wires = Wires(wires)
 
+    # TODO: Remove this conditional branch when using `Counts.process` in devices
     if all_outcomes:
-        return MeasurementProcess(AllCounts, obs=op, wires=wires)
+        return _Counts(AllCounts, obs=op, wires=wires, all_outcomes=all_outcomes)
 
-    return MeasurementProcess(Counts, obs=op, wires=wires)
+    return _Counts(Counts, obs=op, wires=wires, all_outcomes=all_outcomes)
+
+
+# TODO: Make public when removing the ObservableReturnTypes enum
+class _Counts(SampleMeasurement):
+    """Measurement process that returns the samples of a given observable."""
+
+    def __init__(
+        self,
+        return_type: ObservableReturnTypes,
+        obs: Union[Observable, None] = None,
+        wires=None,
+        eigvals=None,
+        id=None,
+        all_outcomes=False,
+    ):
+        self.all_outcomes = all_outcomes
+        super().__init__(return_type, obs, wires, eigvals, id)
+
+    def process_samples(
+        self,
+        samples: Sequence[complex],
+        wire_order: Wires,
+        shot_range: Tuple[int] = None,
+        bin_size: int = None,
+    ):
+        samples = qml.sample(op=self.obs, wires=self._wires).process_samples(
+            samples, wire_order, shot_range, bin_size
+        )
+
+        if bin_size is None:
+            return self._samples_to_counts(samples)
+
+        num_wires = len(self.wires) if self.wires else len(wire_order)
+        shape = (-1, bin_size, num_wires) if self.obs is None else (-1, bin_size)
+
+        return [self._samples_to_counts(bin_sample) for bin_sample in samples.reshape(shape)]
+
+    def _samples_to_counts(self, samples):
+        """Groups the samples into a dictionary showing number of occurences for
+        each possible outcome.
+
+        The format of the dictionary depends on obs.return_type, which is set when
+        calling measurements.counts by setting the kwarg all_outcomes (bool). By default,
+        the dictionary will only contain the observed outcomes. Optionally (all_outcomes=True)
+        the dictionary will instead contain all possible outcomes, with a count of 0
+        for those not observed. See example.
+
+
+        Args:
+            samples: samples in an array of dimension ``(shots,len(wires))``
+
+        Returns:
+            dict: dictionary with format ``{'outcome': num_occurences}``, including all
+                outcomes for the sampled observable
+
+        **Example**
+
+            >>> samples
+            tensor([[0, 0],
+                    [0, 0],
+                    [1, 0]], requires_grad=True)
+
+            By default, this will return:
+            >>> self._samples_to_counts(samples, obs, num_wires)
+            {'00': 2, '10': 1}
+
+            However, if obs.return_type is AllCounts, this will return:
+            >>> self._samples_to_counts(samples, obs, num_wires)
+            {'00': 2, '01': 0, '10': 1, '11': 0}
+
+            The variable all_outcomes can be set when running measurements.counts, i.e.:
+
+             .. code-block:: python3
+
+                dev = qml.device("default.qubit", wires=2, shots=4)
+
+                @qml.qnode(dev)
+                def circuit(x):
+                    qml.RX(x, wires=0)
+                    return qml.counts(all_outcomes=True)
+
+        """
+
+        outcomes = []
+
+        if self.obs is None:
+            # convert samples and outcomes (if using) from arrays to str for dict keys
+            num_wires = len(self.wires) if len(self.wires) > 0 else qml.math.shape(samples)[-1]
+            samples = ["".join([str(s.item()) for s in sample]) for sample in samples]
+            if self.all_outcomes:
+                outcomes = qml.QubitDevice.generate_basis_states(num_wires)
+                outcomes = ["".join([str(o.item()) for o in outcome]) for outcome in outcomes]
+        elif self.all_outcomes:
+            outcomes = qml.eigvals(self.obs)
+
+        # generate empty outcome dict, populate values with state counts
+        outcome_dict = {k: qml.math.int64(0) for k in outcomes}
+        states, _counts = qml.math.unique(samples, return_counts=True)
+        for s, c in zip(states, _counts):
+            outcome_dict[s] = c
+
+        return outcome_dict
+
+    def __copy__(self):
+        return self.__class__(
+            self.return_type,
+            obs=copy.copy(self.obs),
+            eigvals=self._eigvals,
+            wires=self._wires,
+            all_outcomes=self.all_outcomes,
+        )


Test patch:
diff --git a/tests/measurements/test_counts.py b/tests/measurements/test_counts.py
index 0b99e9f4d..947a31954 100644
--- a/tests/measurements/test_counts.py
+++ b/tests/measurements/test_counts.py
@@ -16,7 +16,34 @@ import numpy as np
 import pytest
 
 import pennylane as qml
-from pennylane.measurements import Counts
+from pennylane.measurements import AllCounts, Counts
+from pennylane.operation import Operator
+
+
+# TODO: Remove this when new CustomMP are the default
+def custom_measurement_process(device, spy):
+    assert len(spy.call_args_list) > 0  # make sure method is mocked properly
+
+    samples = device._samples
+    call_args_list = list(spy.call_args_list)
+    for call_args in call_args_list:
+        if not call_args.kwargs.get("counts", False):
+            continue
+        meas = call_args.args[1]
+        shot_range, bin_size = (call_args.kwargs["shot_range"], call_args.kwargs["bin_size"])
+        if isinstance(meas, Operator):
+            all_outcomes = meas.return_type is AllCounts
+            meas = qml.counts(op=meas, all_outcomes=all_outcomes)
+        old_res = device.sample(call_args.args[1], **call_args.kwargs)
+        new_res = meas.process_samples(
+            samples=samples, wire_order=device.wires, shot_range=shot_range, bin_size=bin_size
+        )
+        if isinstance(old_res, dict):
+            old_res = [old_res]
+            new_res = [new_res]
+        for old, new in zip(old_res, new_res):
+            assert old.keys() == new.keys()
+            assert qml.math.allequal(list(old.values()), list(new.values()))
 
 
 class TestCounts:
@@ -39,10 +66,11 @@ class TestCounts:
         ):
             _ = circuit()
 
-    def test_not_an_observable(self):
+    def test_not_an_observable(self, mocker):
         """Test that a UserWarning is raised if the provided
         argument might not be hermitian."""
         dev = qml.device("default.qubit", wires=2, shots=10)
+        spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
         def circuit():
@@ -52,11 +80,14 @@ class TestCounts:
         with pytest.warns(UserWarning, match="Prod might not be hermitian."):
             _ = circuit()
 
-    def test_counts_dimension(self):
+        custom_measurement_process(dev, spy)
+
+    def test_counts_dimension(self, mocker):
         """Test that the counts function outputs counts of the right size"""
         n_sample = 10
 
         dev = qml.device("default.qubit", wires=2, shots=n_sample)
+        spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
         def circuit():
@@ -68,11 +99,14 @@ class TestCounts:
         assert len(sample) == 2
         assert np.all([sum(s.values()) == n_sample for s in sample])
 
-    def test_counts_combination(self):
+        custom_measurement_process(dev, spy)
+
+    def test_counts_combination(self, mocker):
         """Test the output of combining expval, var and counts"""
         n_sample = 10
 
         dev = qml.device("default.qubit", wires=3, shots=n_sample)
+        spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev, diff_method="parameter-shift")
         def circuit():
@@ -89,11 +123,14 @@ class TestCounts:
         assert len(result) == 3
         assert sum(result[0].values()) == n_sample
 
-    def test_single_wire_counts(self):
+        custom_measurement_process(dev, spy)
+
+    def test_single_wire_counts(self, mocker):
         """Test the return type and shape of sampling counts from a single wire"""
         n_sample = 10
 
         dev = qml.device("default.qubit", wires=1, shots=n_sample)
+        spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
         def circuit():
@@ -106,12 +143,15 @@ class TestCounts:
         assert isinstance(result, dict)
         assert sum(result.values()) == n_sample
 
-    def test_multi_wire_counts_regular_shape(self):
+        custom_measurement_process(dev, spy)
+
+    def test_multi_wire_counts_regular_shape(self, mocker):
         """Test the return type and shape of sampling multiple wires
         where a rectangular array is expected"""
         n_sample = 10
 
         dev = qml.device("default.qubit", wires=3, shots=n_sample)
+        spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
         def circuit():
@@ -129,10 +169,13 @@ class TestCounts:
         assert all(sum(r.values()) == n_sample for r in result)
         assert all(all(v.dtype == np.dtype("int") for v in r.values()) for r in result)
 
-    def test_observable_return_type_is_counts(self):
+        custom_measurement_process(dev, spy)
+
+    def test_observable_return_type_is_counts(self, mocker):
         """Test that the return type of the observable is :attr:`ObservableReturnTypes.Counts`"""
         n_shots = 10
         dev = qml.device("default.qubit", wires=1, shots=n_shots)
+        spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
         def circuit():
@@ -142,9 +185,10 @@ class TestCounts:
         circuit()
         assert circuit._qfunc_output.return_type is Counts  # pylint: disable=protected-access
 
-    def test_providing_no_observable_and_no_wires_counts(self):
+    def test_providing_no_observable_and_no_wires_counts(self, mocker):
         """Test that we can provide no observable and no wires to sample function"""
         dev = qml.device("default.qubit", wires=2, shots=1000)
+        spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
         def circuit():
@@ -156,11 +200,14 @@ class TestCounts:
 
         circuit()
 
-    def test_providing_no_observable_and_wires_counts(self):
+        custom_measurement_process(dev, spy)
+
+    def test_providing_no_observable_and_wires_counts(self, mocker):
         """Test that we can provide no observable but specify wires to the sample function"""
         wires = [0, 2]
         wires_obj = qml.wires.Wires(wires)
         dev = qml.device("default.qubit", wires=3, shots=1000)
+        spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
         def circuit():
@@ -173,14 +220,17 @@ class TestCounts:
 
         circuit()
 
+        custom_measurement_process(dev, spy)
+
     @pytest.mark.all_interfaces
     @pytest.mark.parametrize("wires, basis_state", [(None, "010"), ([2, 1], "01")])
     @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
-    def test_counts_no_op_finite_shots(self, interface, wires, basis_state):
+    def test_counts_no_op_finite_shots(self, interface, wires, basis_state, mocker):
         """Check all interfaces with computational basis state counts and
         finite shot"""
         n_shots = 10
         dev = qml.device("default.qubit", wires=3, shots=n_shots)
+        spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev, interface=interface)
         def circuit():
@@ -190,13 +240,16 @@ class TestCounts:
         res = circuit()
         assert res == {basis_state: n_shots}
 
+        custom_measurement_process(dev, spy)
+
     @pytest.mark.all_interfaces
     @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
-    def test_counts_operator_finite_shots(self, interface):
+    def test_counts_operator_finite_shots(self, interface, mocker):
         """Check all interfaces with observable measurement counts and finite
         shot"""
         n_shots = 10
         dev = qml.device("default.qubit", wires=3, shots=n_shots)
+        spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev, interface=interface)
         def circuit():
@@ -205,14 +258,17 @@ class TestCounts:
         res = circuit()
         assert res == {1: n_shots}
 
+        custom_measurement_process(dev, spy)
+
     @pytest.mark.all_interfaces
     @pytest.mark.parametrize("shot_vec", [(1, 10, 10), (1, 10, 1000)])
     @pytest.mark.parametrize("wires, basis_state", [(None, "010"), ([2, 1], "01")])
     @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
-    def test_counts_binned(self, shot_vec, interface, wires, basis_state):
+    def test_counts_binned(self, shot_vec, interface, wires, basis_state, mocker):
         """Check all interfaces with computational basis state counts and
         different shot vectors"""
         dev = qml.device("default.qubit", wires=3, shots=shot_vec)
+        spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev, interface=interface)
         def circuit():
@@ -228,13 +284,16 @@ class TestCounts:
         assert len(res) == len(shot_vec)
         assert sum(sum(res_bin.values()) for res_bin in res) == sum(shot_vec)
 
+        custom_measurement_process(dev, spy)
+
     @pytest.mark.all_interfaces
     @pytest.mark.parametrize("shot_vec", [(1, 10, 10), (1, 10, 1000)])
     @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
-    def test_counts_operator_binned(self, shot_vec, interface):
+    def test_counts_operator_binned(self, shot_vec, interface, mocker):
         """Check all interfaces with observable measurement counts and different
         shot vectors"""
         dev = qml.device("default.qubit", wires=3, shots=shot_vec)
+        spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev, interface=interface)
         def circuit():
@@ -249,11 +308,14 @@ class TestCounts:
         assert len(res) == len(shot_vec)
         assert sum(sum(res_bin.values()) for res_bin in res) == sum(shot_vec)
 
+        custom_measurement_process(dev, spy)
+
     @pytest.mark.parametrize("shot_vec", [(1, 10, 10), (1, 10, 1000)])
-    def test_counts_binned_4_wires(self, shot_vec):
+    def test_counts_binned_4_wires(self, shot_vec, mocker):
         """Check the autograd interface with computational basis state counts and
         different shot vectors on a device with 4 wires"""
         dev = qml.device("default.qubit", wires=4, shots=shot_vec)
+        spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev, interface="autograd")
         def circuit():
@@ -272,11 +334,14 @@ class TestCounts:
         assert len(res) == len(shot_vec)
         assert sum(sum(res_bin.values()) for res_bin in res) == sum(shot_vec)
 
+        custom_measurement_process(dev, spy)
+
     @pytest.mark.parametrize("shot_vec", [(1, 10, 10), (1, 10, 1000)])
-    def test_counts_operator_binned_4_wires(self, shot_vec):
+    def test_counts_operator_binned_4_wires(self, shot_vec, mocker):
         """Check the autograd interface with observable samples to obtain
         counts from and different shot vectors on a device with 4 wires"""
         dev = qml.device("default.qubit", wires=4, shots=shot_vec)
+        spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev, interface="autograd")
         def circuit():
@@ -295,6 +360,8 @@ class TestCounts:
         assert len(res) == len(shot_vec)
         assert sum(sum(res_bin.values()) for res_bin in res) == sum(shot_vec)
 
+        custom_measurement_process(dev, spy)
+
     meas2 = [
         qml.expval(qml.PauliZ(0)),
         qml.var(qml.PauliZ(0)),
@@ -307,10 +374,9 @@ class TestCounts:
     @pytest.mark.parametrize("meas2", meas2)
     @pytest.mark.parametrize("shots", [1000, (1, 10)])
     @pytest.mark.filterwarnings("ignore:Creating an ndarray from ragged nested sequences")
-    def test_counts_observable_finite_shots(self, interface, meas2, shots):
+    def test_counts_observable_finite_shots(self, interface, meas2, shots, mocker):
         """Check all interfaces with observable measurement counts and finite
         shot"""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
 
         if interface == "jax" and meas2.return_type in (
             qml.measurements.Probability,
@@ -319,6 +385,9 @@ class TestCounts:
             reason = "Using the JAX interface, sample and probability measurements cannot be mixed with other measurement types."
             pytest.skip(reason)
 
+        dev = qml.device("default.qubit", wires=3, shots=shots)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
         @qml.qnode(dev, interface=interface)
         def circuit():
             qml.PauliX(0)
@@ -332,12 +401,15 @@ class TestCounts:
         for ind in counts_term_indices:
             assert isinstance(res[ind], dict)
 
-    def test_all_outcomes_kwarg_providing_observable(self):
+        custom_measurement_process(dev, spy)
+
+    def test_all_outcomes_kwarg_providing_observable(self, mocker):
         """Test that the dictionary keys *all* eigenvalues of the observable,
         including 0 count values, if observable is given and all_outcomes=True"""
 
         n_shots = 10
         dev = qml.device("default.qubit", wires=1, shots=n_shots)
+        spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
         def circuit():
@@ -348,13 +420,16 @@ class TestCounts:
 
         assert res == {1: n_shots, -1: 0}
 
-    def test_all_outcomes_kwarg_no_observable_no_wires(self):
+        custom_measurement_process(dev, spy)
+
+    def test_all_outcomes_kwarg_no_observable_no_wires(self, mocker):
         """Test that the dictionary keys are *all* the possible combinations
         of basis states for the device, including 0 count values, if no wire
         count and no observable are given and all_outcomes=True"""
 
         n_shots = 10
         dev = qml.device("default.qubit", wires=2, shots=n_shots)
+        spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
         def circuit():
@@ -364,13 +439,16 @@ class TestCounts:
 
         assert res == {"00": n_shots, "01": 0, "10": 0, "11": 0}
 
-    def test_all_outcomes_kwarg_providing_wires_and_no_observable(self):
+        custom_measurement_process(dev, spy)
+
+    def test_all_outcomes_kwarg_providing_wires_and_no_observable(self, mocker):
         """Test that the dictionary keys are *all* possible combinations
         of basis states for the specified wires, including 0 count values,
         if wire count is given and all_outcomes=True"""
 
         n_shots = 10
         dev = qml.device("default.qubit", wires=4, shots=n_shots)
+        spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
         def circuit():
@@ -380,12 +458,16 @@ class TestCounts:
 
         assert res == {"00": n_shots, "01": 0, "10": 0, "11": 0}
 
-    def test_all_outcomes_hermitian(self):
+        custom_measurement_process(dev, spy)
+
+    def test_all_outcomes_hermitian(self, mocker):
         """Tests that the all_outcomes=True option for counts works with the
         qml.Hermitian observable"""
 
         n_shots = 10
         dev = qml.device("default.qubit", wires=2, shots=n_shots)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
         A = np.array([[1, 0], [0, -1]])
 
         @qml.qnode(dev)
@@ -396,11 +478,14 @@ class TestCounts:
 
         assert res == {-1.0: 0, 1.0: n_shots}
 
-    def test_all_outcomes_multiple_measurements(self):
+        custom_measurement_process(dev, spy)
+
+    def test_all_outcomes_multiple_measurements(self, mocker):
         """Tests that the all_outcomes=True option for counts works when
         multiple measurements are performed"""
 
         dev = qml.device("default.qubit", wires=2, shots=10)
+        spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
         def circuit():
@@ -411,3 +496,5 @@ class TestCounts:
         assert len(res[0]) == 10
         assert res[1] == {"00": 10}
         assert res[2] == {"00": 10, "01": 0, "10": 0, "11": 0}
+
+        custom_measurement_process(dev, spy)