Repository: PennyLaneAI/pennylane

Issue:
Add `Shots.bins()` generator method
[`qml.measurements.Shots`](https://github.com/PennyLaneAI/pennylane/blob/c2f5d44ae3428ac3bdc5fa82a776594110e935fb/pennylane/measurements/shots.py#L46) is our class for keeping track of how many shots to use in calculating measurements.

One of the reasons we need a class like this is to better interact with and support *shot vectors*.  A shot vector is when the same measurements as calculated multiple times using subsets of the total shots.  For example, `shots=(1, 10, 100, 100)` requests that measurements are done with `1` shot, `10` shots, `100` shots, and `100` shots again:

```
@qml.qnode(qml.device('default.qubit'))
def circuit():
    return qml.counts(wires=0)

circuit(shots=(1, 10, 100,100))
```
```
({'0': tensor(1, requires_grad=True)},
 {'0': tensor(10, requires_grad=True)},
 {'0': tensor(100, requires_grad=True)},
 {'0': tensor(100, requires_grad=True)})
```

If we calculated 211 total shots, measurement 1 would be calculated from shots `0:1`, measurement 2 from `1:11`, measurement 3 from `11:111`, and 3 from `111:211`.

This task is to add a property that makes it easy and convenient to access these ranges as "bins", a tuple of `(lower_bound, upper_bound)`. 

For example:

```
>>> shots = Shots((1, 10, 100, 100))
>>> shots.bins()
<generator object <genexpr> at 0x3095a5630>
>>> list(shots.bins())
[(0,1), (1,11), (11,111), (111,211)]
>>> for bin in shots.bins():
...     print(bin)
(0,1)
(1,11)
(11,111)
(111,211)
```

To complete this task, the contributor should:

1) Add the `bins` method to `Shots` in `pennylane/measurements/shots.py`
2) Add tests for the new method in [`tests/measurements/test_shots.py`](https://github.com/PennyLaneAI/pennylane/blob/master/tests/measurements/test_shots.py)
3) Add a changelog entry under `Community Contributions`
4) Add your name to the list of contributors at the bottom of the changelog
5) Make sure that all tests pass and the code adheres to our black and pylint standards




Discussion:
Hey @albi3ro, would it be possible for me to work on this issue? Could you please assign it to me?
Thanks for reaching out @Tarun-Kumar07 . I've assigned you the issue.
Hey @albi3ro, I was wondering about the method  `bins`. Is it part of a larger task or project?
More of something that could make things easier in the future.

An example of where this could be used is here:

https://github.com/PennyLaneAI/pennylane/blob/04be31747c2946d1561f35ea9487a9b486778983/pennylane/devices/qubit/sampling.py#L268

Currently in `devices/qubit/sampling.py`, we handle shot vectors by sampling multiple times and processing the shots independently.  In an ideal world, we would want to just sample once, and then process each bin separately. We made compromises due to time constraints and conceptual ease. This method would allow us to move toward the more ideal way of implementing sampling.

So a follow-on to this task would be to improve the implementation of `_measure_with_samples_diagonalizing_gates`.

Hey @albi3ro just wanted to confirm the case when there are duplicates i.e. shot quantity is repeated twice,
is below the expected behaviour
```
>>> shots = Shots((1, 1, 2, 3))
>>> list(shots.bins())
[(0,1), (1,2), (2,4), (4,7)]
```
Exactly 👍 

Another quick note that the same shots class can also be created with:

```
>>> qml.measurements.Shots(((1,2), 2, 3))
Shots(total_shots=7, shot_vector=(ShotCopies(1 shots x 2), ShotCopies(2 shots x 1), ShotCopies(3 shots x 1)))
```

Resolving patch:
diff --git a/pennylane/devices/qubit/sampling.py b/pennylane/devices/qubit/sampling.py
--- a/pennylane/devices/qubit/sampling.py
+++ b/pennylane/devices/qubit/sampling.py
@@ -268,31 +268,6 @@ def _process_single_shot(samples):
 
         return tuple(processed)
 
-    # if there is a shot vector, build a list containing results for each shot entry
-    if shots.has_partitioned_shots:
-        processed_samples = []
-        for s in shots:
-            # currently we call sample_state for each shot entry, but it may be
-            # better to call sample_state just once with total_shots, then use
-            # the shot_range keyword argument
-            try:
-                samples = sample_state(
-                    state,
-                    shots=s,
-                    is_state_batched=is_state_batched,
-                    wires=wires,
-                    rng=rng,
-                    prng_key=prng_key,
-                )
-            except ValueError as e:
-                if str(e) != "probabilities contain NaN":
-                    raise e
-                samples = qml.math.full((s, len(wires)), 0)
-
-            processed_samples.append(_process_single_shot(samples))
-
-        return tuple(zip(*processed_samples))
-
     try:
         samples = sample_state(
             state,
@@ -307,7 +282,15 @@ def _process_single_shot(samples):
             raise e
         samples = qml.math.full((shots.total_shots, len(wires)), 0)
 
-    return _process_single_shot(samples)
+    processed_samples = []
+    for lower, upper in shots.bins():
+        shot = _process_single_shot(samples[..., lower:upper, :])
+        processed_samples.append(shot)
+
+    if shots.has_partitioned_shots:
+        return tuple(zip(*processed_samples))
+
+    return processed_samples[0]
 
 
 def _measure_classical_shadow(
diff --git a/pennylane/measurements/shots.py b/pennylane/measurements/shots.py
--- a/pennylane/measurements/shots.py
+++ b/pennylane/measurements/shots.py
@@ -261,3 +261,20 @@ def has_partitioned_shots(self):
     def num_copies(self):
         """The total number of copies of any shot quantity."""
         return sum(s.copies for s in self.shot_vector)
+
+    def bins(self):
+        """
+        Yields:
+            tuple: A tuple containing the lower and upper bounds for each shot quantity in shot_vector.
+
+        Example:
+            >>> shots = Shots((1, 1, 2, 3))
+            >>> list(shots.bins())
+            [(0,1), (1,2), (2,4), (4,7)]
+        """
+        lower_bound = 0
+        for sc in self.shot_vector:
+            for _ in range(sc.copies):
+                upper_bound = lower_bound + sc.shots
+                yield lower_bound, upper_bound
+                lower_bound = upper_bound


Test patch:
diff --git a/tests/devices/qubit/test_sampling.py b/tests/devices/qubit/test_sampling.py
--- a/tests/devices/qubit/test_sampling.py
+++ b/tests/devices/qubit/test_sampling.py
@@ -856,7 +856,7 @@ def test_nonsample_measure_shot_vector(self, shots, measurement, expected):
             r = r[0]
 
             assert r.shape == expected.shape
-            assert np.allclose(r, expected, atol=0.01)
+            assert np.allclose(r, expected, atol=0.02)
 
 
 @pytest.mark.jax
@@ -1071,7 +1071,7 @@ def test_nonsample_measure_shot_vector(self, mocker, shots, measurement, expecte
             r = r[0]
 
             assert r.shape == expected.shape
-            assert np.allclose(r, expected, atol=0.01)
+            assert np.allclose(r, expected, atol=0.03)
 
 
 class TestHamiltonianSamples:
diff --git a/tests/measurements/test_shots.py b/tests/measurements/test_shots.py
--- a/tests/measurements/test_shots.py
+++ b/tests/measurements/test_shots.py
@@ -283,3 +283,23 @@ def test_shots_rmul(self):
         scaled_sh1 = 2 * sh1
         rev_scaled_sh1 = sh1 * 2
         assert scaled_sh1.total_shots == rev_scaled_sh1.total_shots
+
+
+class TestShotsBins:
+    """Tests Shots.bins() method."""
+
+    def test_when_shots_is_none(self):
+        """Tests that the method returns no bins when shots is None."""
+        shots = Shots(None)
+        assert not list(shots.bins())
+
+    def test_when_shots_is_int(self):
+        """Tests that the method returns the correct bins when shots is an int."""
+        shots = Shots(10)
+        assert list(shots.bins()) == [(0, 10)]
+
+    @pytest.mark.parametrize("sequence", [[1, 1, 3, 4], [(1, 2), 3, 4]])
+    def test_when_shots_is_sequence_with_copies(self, sequence):
+        """Tests that the method returns the correct bins when shots is a sequence with copies."""
+        shots = Shots(sequence)
+        assert list(shots.bins()) == [(0, 1), (1, 2), (2, 5), (5, 9)]