Repository: PennyLaneAI/pennylane (python)

Issue:
[BUG] Counts does not work with broadcasting
### Expected behavior

Broadcasted parameters should work with counts.

### Actual behavior

The following code returns:
```
Counts: {'0000011010010010100000000001000000010110000000000000000101001110100011010000000001010111001001100001': 1, '0000111100001000111010000101010001101000010001010001000110010010011000101000010001111000010001010000': 1}
Counts: {'1010011011000010100111010011101000000100001000000000001100000000110000000000100100010011000000010111': 1}
```

### Additional information

-

### Source code

```shell
import pennylane as qml

dev = qml.device('default.qubit', wires=2, shots=100)

@qml.qnode(dev, interface="autograd", diff_method=None)
def circuit(param):
    # These two gates represent our QML model.
    qml.RX(param, wires=0)
    qml.CNOT(wires=[0, 1])
    return qml.counts(wires=0)


counts = circuit(qml.numpy.array([1.123, 1.1]))
print(f"Counts: {counts}")
counts = circuit(qml.numpy.array([1.123]))
print(f"Counts: {counts}")
```


### Tracebacks

_No response_

### System information

```shell
-
```


### Existing GitHub issues

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

Pull request description:
support broadcasting for counts
**Context:**
Counts could not be broadcasted, but it should be able to be.

**Description of the Change:**
Add broadcasting support for `qml.counts`. Some interfaces didn't even support the non-broadcasted counts in the new return system, so I added some one-liners to add that support as well! I only adde support for the new return system, as well as an explicit error for the interfaces that don't automatically support it with the old return system.

**Benefits:**
Users can broadcast parameters and still get counts back!

**Possible Drawbacks:**
People might suddenly think counts is differentiable with this extra support, but that is not the case : (

**Related GitHub Issues:**
Fixes #3841

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 81b354d6b..4a04092bf 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -120,6 +120,9 @@
   are applied to a large (>8 wires) tensorflow state. When that occurs, the logic falls back to the tensordot logic instead.
   [(#3884)](https://github.com/PennyLaneAI/pennylane/pull/3884/)
 
+* Fixed parameter broadcasting support with `qml.counts` in most cases, and introduced explicit errors otherwise.
+  [(#3876)](https://github.com/PennyLaneAI/pennylane/pull/3876)
+
 <h3>Contributors</h3>
 
 This release contains contributions from (in alphabetical order):
diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index a1dd255b8..ee6e90a9d 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -65,6 +65,11 @@ from pennylane.tape import QuantumScript, QuantumTape
 from pennylane.wires import Wires
 
 
+def _sample_to_str(sample):
+    """Converts a bit-array to a string. For example, ``[0, 1]`` would become '01'."""
+    return "".join(map(str, sample))
+
+
 class QubitDevice(Device):
     """Abstract base class for PennyLane qubit devices.
 
@@ -1706,7 +1711,9 @@ class QubitDevice(Device):
 
 
         Args:
-            samples: samples in an array of dimension ``(shots,len(wires))``
+            samples: An array of samples, with the shape being ``(shots,len(wires))`` if an observable
+                is provided, with sample values being an array of 0s or 1s for each wire. Otherwise, it
+                has shape ``(shots,)``, with sample values being scalar eigenvalues of the observable
             obs (Observable): the observable sampled
             num_wires (int): number of wires the sampled observable was performed on
 
@@ -1744,23 +1751,33 @@ class QubitDevice(Device):
 
         outcomes = []
 
+        # if an observable was provided, batched samples will have shape (batch_size, shots)
+        batched_ndims = 2
+        shape = samples.shape
+
         if isinstance(obs, CountsMP):
             # convert samples and outcomes (if using) from arrays to str for dict keys
-            samples = ["".join([str(s.item()) for s in sample]) for sample in samples]
-
+            samples = np.apply_along_axis(_sample_to_str, -1, samples)
+            batched_ndims = 3  # no observable was provided, batched samples will have shape (batch_size, shots, len(wires))
             if obs.all_outcomes:
-                outcomes = self.generate_basis_states(num_wires)
-                outcomes = ["".join([str(o.item()) for o in outcome]) for outcome in outcomes]
+                outcomes = list(map(_sample_to_str, self.generate_basis_states(num_wires)))
         elif obs.return_type is AllCounts:
             outcomes = qml.eigvals(obs)
 
-        # generate empty outcome dict, populate values with state counts
-        outcome_dict = {k: np.int64(0) for k in outcomes}
-        states, counts = np.unique(samples, return_counts=True)
-        for s, c in zip(states, counts):
-            outcome_dict[s] = c
+        batched = len(shape) == batched_ndims
+        if not batched:
+            samples = samples[None]
 
-        return outcome_dict
+        # generate empty outcome dict, populate values with state counts
+        base_dict = {k: np.int64(0) for k in outcomes}
+        outcome_dicts = [base_dict.copy() for _ in range(shape[0])]
+        results = [np.unique(batch, return_counts=True) for batch in samples]
+        for result, outcome_dict in zip(results, outcome_dicts):
+            states, counts = result
+            for state, count in zip(states, counts):
+                outcome_dict[state] = count
+
+        return outcome_dicts if batched else outcome_dicts[0]
 
     def sample(self, observable, shot_range=None, bin_size=None, counts=False):
         """Return samples of an observable.
diff --git a/pennylane/interfaces/jax.py b/pennylane/interfaces/jax.py
index 81d1e18fa..14e5adc4e 100644
--- a/pennylane/interfaces/jax.py
+++ b/pennylane/interfaces/jax.py
@@ -169,7 +169,14 @@ def _execute(
         """Auxiliary function to convert the result of a tape to an array,
         unless the tape had Counts measurements that are represented with
         dictionaries. JAX NumPy arrays don't support dictionaries."""
-        return jnp.array(r) if not any(isinstance(m, CountsMP) for m in tape.measurements) else r
+        if any(isinstance(m, CountsMP) for m in tape.measurements):
+            if tape.batch_size is not None:
+                raise InterfaceUnsupportedError(
+                    "Broadcasted circuits with counts return types are only supported with "
+                    "the new return system. Use qml.enable_return() to turn it on."
+                )
+            return r
+        return jnp.array(r)
 
     @jax.custom_vjp
     def wrapped_exec(params):
@@ -563,17 +570,24 @@ def _copy_tape(t, a):
     return tc
 
 
+def _is_count_result(r):
+    """Checks if ``r`` is a single count (or broadcasted count) result"""
+    return isinstance(r, dict) or isinstance(r, list) and all(isinstance(i, dict) for i in r)
+
+
 def _to_jax(res):
     """From a list of tapes results (each result is either a np.array or tuple), transform it to a list of Jax
     results (structure stay the same)."""
     res_ = []
     for r in res:
-        if not isinstance(r, tuple):
+        if _is_count_result(r):
+            res_.append(r)
+        elif not isinstance(r, tuple):
             res_.append(jnp.array(r))
         else:
             sub_r = []
             for r_i in r:
-                if isinstance(r_i, dict):
+                if _is_count_result(r_i):
                     sub_r.append(r_i)
                 else:
                     sub_r.append(jnp.array(r_i))
@@ -586,5 +600,4 @@ def _to_jax_shot_vector(res):
 
     The expected structure of the inputs is a list of tape results with each element in the list being a tuple due to execution using shot vectors.
     """
-    res_ = [tuple(_to_jax([r_])[0] for r_ in r) for r in res]
-    return res_
+    return [tuple(_to_jax([r_])[0] for r_ in r) for r in res]
diff --git a/pennylane/interfaces/tensorflow.py b/pennylane/interfaces/tensorflow.py
index 64961f59f..1cc95e2f3 100644
--- a/pennylane/interfaces/tensorflow.py
+++ b/pennylane/interfaces/tensorflow.py
@@ -24,6 +24,7 @@ from tensorflow.python.eager import context
 
 import pennylane as qml
 from pennylane._device import _get_num_copies
+from pennylane.interfaces import InterfaceUnsupportedError
 from pennylane.measurements import CountsMP
 
 
@@ -76,10 +77,14 @@ def _to_tensors(x):
     Convert a nested tuple structure of arrays into a nested tuple
     structure of TF tensors
     """
-    if not isinstance(x, tuple):
-        return tf.convert_to_tensor(x)
+    if isinstance(x, dict) or isinstance(x, list) and all(isinstance(i, dict) for i in x):
+        # qml.counts returns a dict (list of dicts when broadcasted), can't form a valid tensor
+        return x
 
-    return tuple(_to_tensors(x_) for x_ in x)
+    if isinstance(x, tuple):
+        return tuple(_to_tensors(x_) for x_ in x)
+
+    return tf.convert_to_tensor(x)
 
 
 def _res_restructured(res, tapes, shots=None):
@@ -100,7 +105,7 @@ def _res_restructured(res, tapes, shots=None):
             shot_res_nested.append(shot_res[0] if num_meas == 1 else shot_res)
             start += num_meas
 
-        res_nested.append(shot_res_nested[0] if not shot_vector else tuple(shot_res_nested))
+        res_nested.append(tuple(shot_res_nested) if shot_vector else shot_res_nested[0])
 
     return tuple(res_nested)
 
@@ -193,6 +198,11 @@ def execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_d
         # convert output to TensorFlow tensors
 
         if any(isinstance(m, CountsMP) for m in tape.measurements):
+            if tape.batch_size is not None:
+                raise InterfaceUnsupportedError(
+                    "Broadcasted circuits with counts return types are only supported with "
+                    "the new return system. Use qml.enable_return() to turn it on."
+                )
             continue
 
         if isinstance(res[i], np.ndarray):
@@ -275,7 +285,7 @@ def execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_d
                     with qml.tape.Unwrap(*tapes, params=params_unwrapped):
                         vjps = _compute_vjp(dy, gradient_fn(tapes, **gradient_kwargs))
 
-            variables = tfkwargs.get("variables", None)
+            variables = tfkwargs.get("variables")
             return (vjps, variables) if variables is not None else vjps
 
         return res, grad_fn
@@ -331,10 +341,8 @@ def _execute_new(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1,
         res, jacs = execute_fn(tapes, **gradient_kwargs)
 
     for i, r in enumerate(res):
-        # skip in the case of counts
-        if not isinstance(r, dict):
-            # convert output to TensorFlow tensors
-            res[i] = _to_tensors(r)
+        # convert output to TensorFlow tensors
+        res[i] = _to_tensors(r)
 
     @tf.custom_gradient
     def _execute(*parameters):  # pylint:disable=unused-argument
@@ -415,7 +423,7 @@ def _execute_new(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1,
             # filter out untrainable parameters if they happen to appear in the vjp
             vjps = [vjp for vjp in vjps if 0 not in qml.math.shape(vjp)]
 
-            variables = tfkwargs.get("variables", None)
+            variables = tfkwargs.get("variables")
             return (vjps, variables) if variables is not None else vjps
 
         return res, grad_fn
diff --git a/pennylane/interfaces/torch.py b/pennylane/interfaces/torch.py
index 0263020fb..6b1bfe1ae 100644
--- a/pennylane/interfaces/torch.py
+++ b/pennylane/interfaces/torch.py
@@ -246,15 +246,15 @@ def execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_d
         tape.trainable_params = qml.math.get_trainable_indices(params)
         parameters.extend(tape.get_parameters())
 
-    kwargs = dict(
-        tapes=tapes,
-        device=device,
-        execute_fn=execute_fn,
-        gradient_fn=gradient_fn,
-        gradient_kwargs=gradient_kwargs,
-        _n=_n,
-        max_diff=max_diff,
-    )
+    kwargs = {
+        "tapes": tapes,
+        "device": device,
+        "execute_fn": execute_fn,
+        "gradient_fn": gradient_fn,
+        "gradient_kwargs": gradient_kwargs,
+        "_n": _n,
+        "max_diff": max_diff,
+    }
     return ExecuteTapes.apply(kwargs, *parameters)
 
 
@@ -487,15 +487,15 @@ def _execute_new(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1,
         tape.trainable_params = qml.math.get_trainable_indices(params)
         parameters.extend(tape.get_parameters())
 
-    kwargs = dict(
-        tapes=tapes,
-        device=device,
-        execute_fn=execute_fn,
-        gradient_fn=gradient_fn,
-        gradient_kwargs=gradient_kwargs,
-        _n=_n,
-        max_diff=max_diff,
-    )
+    kwargs = {
+        "tapes": tapes,
+        "device": device,
+        "execute_fn": execute_fn,
+        "gradient_fn": gradient_fn,
+        "gradient_kwargs": gradient_kwargs,
+        "_n": _n,
+        "max_diff": max_diff,
+    }
 
     return ExecuteTapesNew.apply(kwargs, *parameters)
 
@@ -505,12 +505,15 @@ def _res_to_torch(r, ctx):
     if isinstance(r, (list, tuple)):
         res = []
         for t in r:
-            if isinstance(t, dict):
+            if isinstance(t, dict) or isinstance(t, list) and all(isinstance(i, dict) for i in t):
+                # count result, single or broadcasted
                 res.append(t)
             else:
                 res.append(torch.as_tensor(t, device=ctx.torch_device))
         if isinstance(r, tuple):
             res = tuple(res)
+    elif isinstance(r, dict):
+        res = r
     else:
         res = torch.as_tensor(r, device=ctx.torch_device)
 
@@ -520,13 +523,12 @@ def _res_to_torch(r, ctx):
 def _jac_to_torch(i, ctx):
     """Convert Jacobian from unwrapped execution to torch in the given ctx."""
     if ctx.jacs:
-        jacobians = []
-
         multi_m = len(ctx.tapes[i].measurements) > 1
         multi_p = len(ctx.tapes[i].trainable_params) > 1
 
         # Multiple measurements and parameters: Jacobian is a tuple of tuple
         if multi_p and multi_m:
+            jacobians = []
             for jacobian in ctx.jacs[i]:
                 inside_nested_jacobian = [
                     torch.as_tensor(j, device=ctx.torch_device) for j in jacobian
diff --git a/pennylane/measurements/counts.py b/pennylane/measurements/counts.py
index 12e7e8fca..fb4363fcb 100644
--- a/pennylane/measurements/counts.py
+++ b/pennylane/measurements/counts.py
@@ -25,6 +25,11 @@ from pennylane.wires import Wires
 from .measurements import AllCounts, Counts, SampleMeasurement
 
 
+def _sample_to_str(sample):
+    """Converts a bit-array to a string. For example, ``[0, 1]`` would become '01'."""
+    return "".join(map(str, sample))
+
+
 def counts(op=None, wires=None, all_outcomes=False):
     r"""Sample from the supplied observable, with the number of shots
     determined from the ``dev.shots`` attribute of the corresponding device,
@@ -212,7 +217,9 @@ class CountsMP(SampleMeasurement):
         for those not observed. See example.
 
         Args:
-            samples: samples in an array of dimension ``(shots,len(wires))``
+            samples: An array of samples, with the shape being ``(shots,len(wires))`` if an observable
+                is provided, with sample values being an array of 0s or 1s for each wire. Otherwise, it
+                has shape ``(shots,)``, with sample values being scalar eigenvalues of the observable
 
         Returns:
             dict: dictionary with format ``{'outcome': num_occurrences}``, including all
@@ -248,23 +255,36 @@ class CountsMP(SampleMeasurement):
 
         outcomes = []
 
+        # if an observable was provided, batched samples will have shape (batch_size, shots)
+        batched_ndims = 2
+        shape = qml.math.shape(samples)
+
         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]
+            samples = qml.math.apply_along_axis(_sample_to_str, -1, samples)
+            batched_ndims = 3  # no observable was provided, batched samples will have shape (batch_size, shots, len(wires))
             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]
+                num_wires = len(self.wires) if len(self.wires) > 0 else shape[-1]
+                outcomes = list(
+                    map(_sample_to_str, qml.QubitDevice.generate_basis_states(num_wires))
+                )
         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
+        batched = len(shape) == batched_ndims
+        if not batched:
+            samples = samples[None]
 
-        return outcome_dict
+        # generate empty outcome dict, populate values with state counts
+        base_dict = {k: qml.math.int64(0) for k in outcomes}
+        outcome_dicts = [base_dict.copy() for _ in range(shape[0])]
+        results = [qml.math.unique(batch, return_counts=True) for batch in samples]
+        for result, outcome_dict in zip(results, outcome_dicts):
+            states, _counts = result
+            for state, count in zip(states, _counts):
+                outcome_dict[state] = count
+
+        return outcome_dicts if batched else outcome_dicts[0]
 
     def __copy__(self):
         return self.__class__(


Test patch:
diff --git a/tests/measurements/test_counts.py b/tests/measurements/test_counts.py
index bf5c245e2..8ef31891c 100644
--- a/tests/measurements/test_counts.py
+++ b/tests/measurements/test_counts.py
@@ -18,6 +18,7 @@ import numpy as np
 import pytest
 
 import pennylane as qml
+from pennylane.interfaces import InterfaceUnsupportedError
 from pennylane.measurements import AllCounts, Counts, CountsMP
 from pennylane.operation import Operator
 from pennylane.wires import Wires
@@ -27,7 +28,7 @@ from pennylane.wires import Wires
 def custom_measurement_process(device, spy):
     assert len(spy.call_args_list) > 0  # make sure method is mocked properly
 
-    samples = device._samples
+    samples = device._samples  # pylint:disable=protected-access
     call_args_list = list(spy.call_args_list)
     for call_args in call_args_list:
         if not call_args.kwargs.get("counts", False):
@@ -52,6 +53,8 @@ def custom_measurement_process(device, spy):
 class TestCounts:
     """Tests for the counts function"""
 
+    # pylint:disable=too-many-public-methods
+
     def test_counts_properties(self):
         """Test that the properties are correct."""
         meas1 = qml.counts(wires=0)
@@ -78,7 +81,7 @@ class TestCounts:
         meas = qml.counts(wires=0, all_outcomes=True)
         meas_copy = copy.copy(meas)
         assert meas_copy.wires == Wires(0)
-        assert meas_copy.all_outcomes == True
+        assert meas_copy.all_outcomes is True
 
     def test_providing_observable_and_wires(self):
         """Test that a ValueError is raised if both an observable is provided and wires are
@@ -132,6 +135,46 @@ class TestCounts:
 
         custom_measurement_process(dev, spy)
 
+    def test_batched_counts_dimension(self, mocker):
+        """Test that the counts function outputs counts of the right size with batching"""
+        n_sample = 10
+
+        dev = qml.device("default.qubit", wires=2, shots=n_sample)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.RX([0.54, 0.65], wires=0)
+            return qml.counts(qml.PauliZ(0)), qml.counts(qml.PauliX(1))
+
+        sample = circuit()
+
+        assert isinstance(sample, tuple)
+        assert len(sample) == 2
+        assert np.all([sum(s.values()) == n_sample for batch in sample for s in batch])
+
+        custom_measurement_process(dev, spy)
+
+    @pytest.mark.parametrize(
+        "interface",
+        [
+            pytest.param("jax", marks=pytest.mark.jax),
+            pytest.param("tensorflow", marks=pytest.mark.tf),
+        ],
+    )
+    def test_batched_counts_not_supported_old_return(self, interface):
+        """Test that counts with batching is not supporting in the old return system."""
+
+        @qml.qnode(qml.device("default.qubit", wires=2, shots=10), interface=interface)
+        def circuit():
+            qml.RX([0.54, 0.65], wires=0)
+            return qml.counts(qml.PauliZ(0))
+
+        with pytest.raises(
+            InterfaceUnsupportedError, match="Broadcasted circuits with counts return types"
+        ):
+            _ = circuit()
+
     def test_counts_combination(self, mocker):
         """Test the output of combining expval, var and counts"""
         n_sample = 10
@@ -202,11 +245,10 @@ class TestCounts:
 
         custom_measurement_process(dev, spy)
 
-    def test_observable_return_type_is_counts(self, mocker):
+    def test_observable_return_type_is_counts(self):
         """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():
@@ -253,6 +295,20 @@ class TestCounts:
 
         custom_measurement_process(dev, spy)
 
+    def test_batched_counts_work_individually(self, mocker):
+        """Test that each counts call operates independently"""
+        n_shots = 10
+        dev = qml.device("default.qubit", wires=1, shots=n_shots)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.pow(qml.PauliX(0), z=[1, 2])
+            return qml.counts()
+
+        assert all(circuit() == [{"1": 10}, {"0": 10}])
+        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"])
@@ -295,7 +351,9 @@ class TestCounts:
     @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, mocker):
+    def test_counts_binned(
+        self, shot_vec, interface, wires, basis_state, mocker
+    ):  # pylint:disable=too-many-arguments
         """Check all interfaces with computational basis state counts and
         different shot vectors"""
         dev = qml.device("default.qubit", wires=3, shots=shot_vec)
@@ -529,6 +587,21 @@ class TestCounts:
         assert res[2] == {"00": 10, "01": 0, "10": 0, "11": 0}
         custom_measurement_process(dev, spy)
 
+    def test_batched_all_outcomes(self, mocker):
+        """Tests that all_outcomes=True works with broadcasting."""
+        n_shots = 10
+        dev = qml.device("default.qubit", wires=1, shots=n_shots)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.pow(qml.PauliX(0), z=[1, 2])
+            return qml.counts(qml.PauliZ(0), all_outcomes=True)
+
+        assert all(circuit() == [{1: 0, -1: n_shots}, {1: n_shots, -1: 0}])
+
+        custom_measurement_process(dev, spy)
+
     def test_counts_empty_wires(self):
         """Test that using ``qml.counts`` with an empty wire list raises an error."""
         with pytest.raises(ValueError, match="Cannot set an empty list of wires."):
diff --git a/tests/returntypes/core/test_new_counts.py b/tests/returntypes/core/test_new_counts.py
new file mode 100644
index 000000000..419d090cf
--- /dev/null
+++ b/tests/returntypes/core/test_new_counts.py
@@ -0,0 +1,147 @@
+# Copyright 2018-2023 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for the qml.counts measurement process."""
+import pytest
+
+import pennylane as qml
+from pennylane.measurements import AllCounts
+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  # pylint:disable=protected-access
+    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()))
+
+
+@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(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():
+        qml.PauliX(1)
+        return qml.counts(wires=wires)
+
+    res = circuit()
+    assert res == {basis_state: n_shots}
+
+    custom_measurement_process(dev, spy)
+
+
+@pytest.mark.all_interfaces
+@pytest.mark.parametrize("wires, basis_states", [(None, ("010", "000")), ([2, 1], ("01", "00"))])
+@pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
+def test_batched_counts_no_op_finite_shots(interface, wires, basis_states, 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():
+        qml.pow(qml.PauliX(1), z=[1, 2])
+        return qml.counts(wires=wires)
+
+    assert circuit() == [{basis_state: n_shots} for basis_state in basis_states]
+
+    custom_measurement_process(dev, spy)
+
+
+@pytest.mark.all_interfaces
+@pytest.mark.parametrize("wires, basis_states", [(None, ("010", "000")), ([2, 1], ("01", "00"))])
+@pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
+def test_batched_counts_and_expval_no_op_finite_shots(interface, wires, basis_states, 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():
+        qml.pow(qml.PauliX(1), z=[1, 2])
+        return qml.counts(wires=wires), qml.expval(qml.PauliZ(0))
+
+    res = circuit()
+    assert isinstance(res, tuple) and len(res) == 2
+    assert res[0] == [{basis_state: n_shots} for basis_state in basis_states]
+    assert len(res[1]) == 2 and qml.math.allequal(res[1], 1)
+
+    custom_measurement_process(dev, spy)
+
+
+@pytest.mark.all_interfaces
+@pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
+def test_batched_counts_operator_finite_shots(interface, mocker):
+    """Check all interfaces with observable measurement counts, batching and finite shots"""
+    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():
+        qml.pow(qml.PauliX(0), z=[1, 2])
+        return qml.counts(qml.PauliZ(0))
+
+    assert circuit() == [{-1: n_shots}, {1: n_shots}]
+
+    custom_measurement_process(dev, spy)
+
+
+@pytest.mark.all_interfaces
+@pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
+def test_batched_counts_and_expval_operator_finite_shots(interface, mocker):
+    """Check all interfaces with observable measurement counts, batching and finite shots"""
+    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():
+        qml.pow(qml.PauliX(0), z=[1, 2])
+        return qml.counts(qml.PauliZ(0)), qml.expval(qml.PauliZ(0))
+
+    res = circuit()
+    assert isinstance(res, tuple) and len(res) == 2
+    assert res[0] == [{-1: n_shots}, {1: n_shots}]
+    assert len(res[1]) == 2 and qml.math.allequal(res[1], [-1, 1])
+
+    custom_measurement_process(dev, spy)
diff --git a/tests/tests_passing_pylint.txt b/tests/tests_passing_pylint.txt
index 09603d8a7..edd4a55b7 100644
--- a/tests/tests_passing_pylint.txt
+++ b/tests/tests_passing_pylint.txt
@@ -3,6 +3,7 @@ tests/devices/test_execution_config.py
 tests/drawer/test_tape_mpl.py
 tests/drawer/test_draw_mpl.py
 tests/drawer/test_drawer_utils.py
+tests/measurements/test_counts.py
 tests/measurements/test_probs.py
 tests/qchem/test_givens_rotations.py
 tests/templates/test_subroutines/test_basis_rotation.py
@@ -34,6 +35,7 @@ tests/pulse/**
 tests/math/test_init.py
 tests/math/test_functions.py
 tests/gradients/test_spsa_gradient.py
+tests/returntypes/core/test_new_counts.py
 tests/returntypes/gradients/test_spsa_gradient_new.py
 tests/returntypes/gradients/test_spsa_gradient_shot_vec_new.py
 tests/returntypes/gradients/test_jvp.py