Repository: PennyLaneAI/pennylane
Pull request #1103 (2021-02-22 17:30:25)

[SHOTS] Allows a list of shots to be specified
**Context:** Expands the functionality of the QNode, by allowing measuring statistics to be course grained over shots in a single evaluation.

**Description of the changes:**

Consider

```pycon
>>> shots_list = [5, 10, 1000]
>>> dev = qml.device("default.qubit", wires=2, analytic=False, shots=shots_list)
```

When QNodes are executed on this device, a single execution of 1015 shots will be submitted.
However, three sets of measurement statistics will be returned; using the first 5 shots,
second set of 10 shots, and final 1000 shots, separately.

For example:

```python
@qml.qnode(dev)
def circuit(x):
    qml.RX(x, wires=0)
    qml.CNOT(wires=[0, 1])
    return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0))
```

Executing this, we will get an output of size `(3, 2)`:

```pycon
>>> circuit(0.5)
[[0.33333333 1.        ]
 [0.2        1.        ]
 [0.012      0.868     ]]
```

The output remains fully differentiable.

**Benefits:**

* More fine-grained control over the shots budget, support for shot-adaptive optimizers
* The logic is relatively well-contained, so it should be easy to update this to follow any UI changes to how shots are specified.

**Drawbacks:** Breaks all plugins that depend on `QubitDevice`, due to two new keyword arguments in `expval`, `var`, and `sample`.

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [.github/CHANGELOG.md](/PennyLaneAI/pennylane/blob/master/.github/CHANGELOG.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.
mariaschuld: Nice! 

One thought, we cannot assume that other devices inheriting from  `Device` support sequences of shots. What will happen if this is attempted in those devices?



 
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1103?src=pr&el=h1) Report
> Merging [#1103](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1103?src=pr&el=desc) (d660f9e) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/c8673eb45d9b3d435867f04b9943a5c4b53322be?el=desc) (c8673eb) will **decrease** coverage by `0.31%`.
> The diff coverage is `50.66%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1103/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1103?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master    #1103      +/-   ##
==========================================
- Coverage   97.73%   97.41%   -0.32%     
==========================================
  Files         154      154              
  Lines       11681    11730      +49     
==========================================
+ Hits        11416    11427      +11     
- Misses        265      303      +38     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1103?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [pennylane/\_device.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1103/diff?src=pr&el=tree#diff-cGVubnlsYW5lL19kZXZpY2UucHk=) | `90.62% <33.33%> (-5.55%)` | :arrow_down: |
| [pennylane/\_qubit\_device.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1103/diff?src=pr&el=tree#diff-cGVubnlsYW5lL19xdWJpdF9kZXZpY2UucHk=) | `90.03% <56.14%> (-8.75%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1103?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1103?src=pr&el=footer). Last update [c8673eb...d660f9e](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1103?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

co9olguy: Is `analytic=False` still needed in the example?
josh146: > Is analytic=False still needed in the example?

@co9olguy yes, until #1079 is merged, which removes the `analytic` keyword.
josh146: > One thought, we cannot assume that other devices inheriting from Device support sequences of shots. What will happen if this is attempted in those devices?

For devices that don't know about `self._shot_vector`, nothing changes - they will continue to just use the total number of shots. Don't know if this makes sense or not - we could add something to the capabilities dictionary?
josh146: > I'm not feeling super keen on this right now. It feels like we're adding a ton of complexity for what amounts to a for loop in user code.

Unfortunately, a for loop in the user code (while giving the same results) will generate separate jobs that may be held up in a remote queue. The idea here is to 'batch' shots on a single QNode eval.

Shot adaptive quantum optimizers typically need to generate many 1-shot expectation values, which is currently not possible to implement in PennyLane without submitting many single-shot jobs
josh146: @Thenerdstation have standardized the format to a namedtuple, and made every element a tuple - it cleans it up a _lot_, and helps remove a lot of redundant lines.
josh146: > This feels a lot cleaner thank you!

Np!

> Approval for code structure and readability, but I leave it to nathan and maria to decide if we want to break the existing plugins for this.

I expect #1079 to pre-break all plugins, so took a little liberty here 😆 
Thenerdstation: > I expect #1079 to pre-break all plugins, so took a little liberty here

That's fair

Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index 63cc5c1c6..9392c0d00 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -2,6 +2,42 @@
 
 <h3>New features since last release</h3>
 
+* Batches of shots can now be specified as a list, allowing measurement statistics
+  to be course-grained with a single QNode evaluation.
+  [(#1103)](https://github.com/PennyLaneAI/pennylane/pull/1103)
+
+  Consider
+
+  ```pycon
+  >>> shots_list = [5, 10, 1000]
+  >>> dev = qml.device("default.qubit", wires=2, analytic=False, shots=shots_list)
+  ```
+
+  When QNodes are executed on this device, a single execution of 1015 shots will be submitted.
+  However, three sets of measurement statistics will be returned; using the first 5 shots,
+  second set of 10 shots, and final 1000 shots, separately.
+
+  For example:
+
+  ```python
+  @qml.qnode(dev)
+  def circuit(x):
+      qml.RX(x, wires=0)
+      qml.CNOT(wires=[0, 1])
+      return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0))
+  ```
+
+  Executing this, we will get an output of size `(3, 2)`:
+
+  ```pycon
+  >>> circuit(0.5)
+  [[0.33333333 1.        ]
+   [0.2        1.        ]
+   [0.012      0.868     ]]
+  ```
+
+  This output remains fully differentiable.
+
 - The number of shots can now be specified on a temporary basis when evaluating a QNode.
   [(#1075)](https://github.com/PennyLaneAI/pennylane/pull/1075)
 
@@ -93,6 +129,34 @@
   argument to change the number of shots on a per-call basis.
   [(#1075)](https://github.com/PennyLaneAI/pennylane/pull/1075)
 
+* For devices inheriting from `QubitDevice`, the methods `expval`, `var`, `sample`
+  accept two new keyword arguments --- `shot_range` and `bin_size`.
+  [(#1103)](https://github.com/PennyLaneAI/pennylane/pull/1103)
+
+  These new arguments allow for the statistics to be performed on only a subset of device samples.
+  This finer level of control is accessible from the main UI by instantiating a device with a batch
+  of shots.
+
+  For example, consider the following device:
+
+  ```pycon
+  >>> dev = qml.device("my_device", shots=[5, (10, 3), 100])
+  ```
+
+  This device will execute QNodes using 135 shots, however
+  measurement statistics will be **course grained** across these 135
+  shots:
+
+  * All measurement statistics will first be computed using the
+    first 5 shots --- that is, `shots_range=[0, 5]`, `bin_size=5`.
+
+  * Next, the tuple `(10, 3)` indicates 10 shots, repeated 3 times. We will want to use
+    `shot_range=[5, 35]`, performing the expectation value in bins of size 10
+    (`bin_size=10`).
+
+  * Finally, we repeat the measurement statistics for the final 100 shots,
+    `shot_range=[35, 135]`, `bin_size=100`.
+
 <h3>Bug fixes</h3>
 
 * The `ExpvalCost` class raises an error if instantiated
diff --git a/doc/introduction/circuits.rst b/doc/introduction/circuits.rst
index 9cfd494be..50e67b242 100644
--- a/doc/introduction/circuits.rst
+++ b/doc/introduction/circuits.rst
@@ -120,6 +120,38 @@ devices, the options are:
 
 For a plugin device, refer to the plugin documentation for available device options.
 
+Shot batches
+^^^^^^^^^^^^
+
+Batches of shots can be specified by passing a list, allowing measurement statistics
+to be course-grained with a single QNode evaluation.
+
+Consider
+
+>>> shots_list = [5, 10, 1000]
+>>> dev = qml.device("default.qubit", wires=2, analytic=False, shots=shots_list)
+
+When QNodes are executed on this device, a single execution of 1015 shots will be submitted.
+However, three sets of measurement statistics will be returned; using the first 5 shots,
+second set of 10 shots, and final 1000 shots, separately.
+
+For example:
+
+.. code-block:: python
+
+    @qml.qnode(dev)
+    def circuit(x):
+      qml.RX(x, wires=0)
+      qml.CNOT(wires=[0, 1])
+      return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0))
+
+Executing this, we will get an output of size ``(3, 2)``:
+
+>>> circuit(0.5)
+[[0.33333333 1.        ]
+[0.2        1.        ]
+[0.012      0.868     ]]
+
 Custom wire labels
 ^^^^^^^^^^^^^^^^^^
 
diff --git a/pennylane/_device.py b/pennylane/_device.py
index 30936169b..e70873614 100644
--- a/pennylane/_device.py
+++ b/pennylane/_device.py
@@ -16,8 +16,8 @@ This module contains the :class:`Device` abstract base class.
 """
 # pylint: disable=too-many-format-args
 import abc
-from collections.abc import Iterable
-from collections import OrderedDict
+from collections.abc import Iterable, Sequence
+from collections import OrderedDict, namedtuple
 
 import numpy as np
 
@@ -36,6 +36,62 @@ from pennylane.qnodes import QuantumFunctionError
 from pennylane.wires import Wires, WireError
 
 
+ShotTuple = namedtuple("ShotTuple", ["shots", "copies"])
+"""tuple[int, int]: Represents copies of a shot number."""
+
+
+def _process_shot_sequence(shot_list):
+    """Process the shot sequence, to determine the total
+    number of shots and the shot vector.
+
+    Args:
+        shot_list (Sequence[int, tuple[int]]): sequence of non-negative shot integers
+
+    Returns:
+        tuple[int, list[.ShotTuple[int]]]: A tuple containing the total number
+        of shots, as well as a list of shot tuples.
+
+    **Example**
+
+    >>> shot_list = [3, 1, 2, 2, 2, 2, 6, 1, 1, 5, 12, 10, 10]
+    >>> _process_shot_sequence(shot_list)
+    (57,
+     [ShotTuple(shots=3, copies=1),
+      ShotTuple(shots=1, copies=1),
+      ShotTuple(shots=2, copies=4),
+      ShotTuple(shots=6, copies=1),
+      ShotTuple(shots=1, copies=2),
+      ShotTuple(shots=5, copies=1),
+      ShotTuple(shots=12, copies=1),
+      ShotTuple(shots=10, copies=2)])
+
+    The total number of shots (57), and a sparse representation of the shot
+    sequence is returned, where tuples indicate the number of times a shot
+    integer is repeated.
+    """
+    if all(isinstance(s, int) for s in shot_list):
+
+        if len(set(shot_list)) == 1:
+            # All shots are identical, only require a single shot tuple
+            shot_vector = [ShotTuple(shots=shot_list[0], copies=len(shot_list))]
+        else:
+            # Iterate through the shots, and group consecutive identical shots
+            split_at_repeated = np.split(shot_list, np.diff(shot_list).nonzero()[0] + 1)
+            shot_vector = [ShotTuple(shots=i[0], copies=len(i)) for i in split_at_repeated]
+
+    elif all(isinstance(s, (int, tuple)) for s in shot_list):
+        # shot list contains tuples; assume it is already in a sparse representation
+        shot_vector = [
+            ShotTuple(*i) if isinstance(i, tuple) else ShotTuple(i, 1) for i in shot_list
+        ]
+
+    else:
+        raise ValueError(f"Unknown shot sequence format {shot_list}")
+
+    total_shots = np.sum(np.prod(shot_vector, axis=1))
+    return total_shots, shot_vector
+
+
 class DeviceError(Exception):
     """Exception raised by a :class:`~.pennylane._device.Device` when it encounters an illegal
     operation in the quantum circuit.
@@ -178,13 +234,23 @@ class Device(abc.ABC):
         Raises:
             DeviceError: if number of shots is less than 1
         """
-        if shots < 1:
+        if isinstance(shots, int):
+            if shots < 1:
+                raise DeviceError(
+                    "The specified number of shots needs to be at least 1. Got {}.".format(shots)
+                )
+
+            self._shots = int(shots)
+            self._shot_vector = None
+
+        elif isinstance(shots, Sequence) and not isinstance(shots, str):
+            self._shots, self._shot_vector = _process_shot_sequence(shots)
+
+        else:
             raise DeviceError(
-                "The specified number of shots needs to be at least 1. Got {}.".format(shots)
+                "Shots must be a single non-negative integer or a sequence of non-negative integers."
             )
 
-        self._shots = int(shots)
-
     def define_wire_map(self, wires):
         """Create the map from user-provided wire labels to the wire labels used by the device.
 
diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index 3455b02c1..56cad73fc 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -17,7 +17,7 @@ This module contains the :class:`QubitDevice` abstract base class.
 
 # For now, arguments may be different from the signatures provided in Device
 # e.g. instead of expval(self, observable, wires, par) have expval(self, observable)
-# pylint: disable=arguments-differ, abstract-method, no-value-for-parameter,too-many-instance-attributes
+# pylint: disable=arguments-differ, abstract-method, no-value-for-parameter,too-many-instance-attributes,too-many-branches
 import abc
 from collections import OrderedDict
 import itertools
@@ -71,8 +71,9 @@ class QubitDevice(Device):
         wires (int, Iterable[Number, str]]): Number of subsystems represented by the device,
             or iterable that contains unique labels for the subsystems as numbers (i.e., ``[-1, 0, 2]``)
             or strings (``['ancilla', 'q1', 'q2']``). Default 1 if not specified.
-        shots (int): number of circuit evaluations/random samples used to estimate
-            expectation values of observables
+        shots (int, list[int]): Number of circuit evaluations/random samples used to estimate
+            expectation values of observables. If a list of integers is passed, the circuit
+            evaluations are batched over the list of shots.
         analytic (bool): If ``True``, the device calculates probability, expectation values,
             and variances analytically. If ``False``, a finite number of samples set by
             the argument ``shots`` are used to estimate these quantities.
@@ -208,7 +209,29 @@ class QubitDevice(Device):
             self._samples = self.generate_samples()
 
         # compute the required statistics
-        results = self.statistics(circuit.observables)
+        if not self.analytic and self._shot_vector is not None:
+
+            results = []
+            s1 = 0
+
+            for shot_tuple in self._shot_vector:
+                s2 = s1 + np.prod(shot_tuple)
+                r = self.statistics(
+                    circuit.observables, shot_range=[s1, s2], bin_size=shot_tuple.shots
+                )
+                r = qml.math.squeeze(r)
+
+                if shot_tuple.copies > 1:
+                    results.extend(r.T)
+                else:
+                    results.append(r)
+
+                s1 = s2
+
+            results = qml.math.stack(results)
+
+        else:
+            results = self.statistics(circuit.observables)
 
         if circuit.all_sampled or not circuit.is_sampled:
             results = self._asarray(results)
@@ -317,7 +340,7 @@ class QubitDevice(Device):
 
         return Wires.all_wires(list_of_wires)
 
-    def statistics(self, observables):
+    def statistics(self, observables, shot_range=None, bin_size=None):
         """Process measurement results from circuit execution and return statistics.
 
         This includes returning expectation values, variance, samples, probabilities, states, and
@@ -325,28 +348,60 @@ class QubitDevice(Device):
 
         Args:
             observables (List[.Observable]): the observables to be measured
+            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:
             QuantumFunctionError: if the value of :attr:`~.Observable.return_type` is not supported
 
         Returns:
             Union[float, List[float]]: the corresponding statistics
+
+        .. UsageDetails::
+
+            The ``shot_range`` and ``bin_size`` arguments allow for the statistics
+            to be performed on only a subset of device samples. This finer level
+            of control is accessible from the main UI by instantiating a device
+            with a batch of shots.
+
+            For example, consider the following device:
+
+            >>> dev = qml.device("my_device", shots=[5, (10, 3), 100])
+
+            This device will execute QNodes using 135 shots, however
+            measurement statistics will be **course grained** across these 135
+            shots:
+
+            * All measurement statistics will first be computed using the
+              first 5 shots --- that is, ``shots_range=[0, 5]``, ``bin_size=5``.
+
+            * Next, the tuple ``(10, 3)`` indicates 10 shots, repeated 3 times. We will want to use
+              ``shot_range=[5, 35]``, performing the expectation value in bins of size 10
+              (``bin_size=10``).
+
+            * Finally, we repeat the measurement statistics for the final 100 shots,
+              ``shot_range=[35, 135]``, ``bin_size=100``.
         """
         results = []
 
         for obs in observables:
             # Pass instances directly
             if obs.return_type is Expectation:
-                results.append(self.expval(obs))
+                results.append(self.expval(obs, shot_range=shot_range, bin_size=bin_size))
 
             elif obs.return_type is Variance:
-                results.append(self.var(obs))
+                results.append(self.var(obs, shot_range=shot_range, bin_size=bin_size))
 
             elif obs.return_type is Sample:
-                results.append(self.sample(obs))
+                results.append(self.sample(obs, shot_range=shot_range, bin_size=bin_size))
 
             elif obs.return_type is Probability:
-                results.append(self.probability(wires=obs.wires))
+                results.append(
+                    self.probability(wires=obs.wires, shot_range=shot_range, bin_size=bin_size)
+                )
 
             elif obs.return_type is State:
                 if len(observables) > 1:
@@ -425,6 +480,7 @@ class QubitDevice(Device):
 
         Args:
             number_of_states (int): the number of basis states to sample from
+            state_probability (array[float]): the computational basis probability vector
 
         Returns:
             List[int]: the sampled basis states
@@ -543,13 +599,18 @@ class QubitDevice(Device):
         """
         raise NotImplementedError
 
-    def estimate_probability(self, wires=None):
+    def estimate_probability(self, wires=None, shot_range=None, bin_size=None):
         """Return the estimated probability of each computational basis state
         using the generated samples.
 
         Args:
             wires (Iterable[Number, str], Number, str, Wires): wires to calculate
                 marginal probabilities for. Wires not provided are traced out of the system.
+            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.
 
         Returns:
             List[float]: list of the probabilities
@@ -561,19 +622,33 @@ class QubitDevice(Device):
         # translate to wire labels used by device
         device_wires = self.map_wires(wires)
 
-        samples = self._samples[:, device_wires]
+        sample_slice = Ellipsis if shot_range is None else slice(*shot_range)
+        samples = self._samples[sample_slice, device_wires]
 
         # convert samples from a list of 0, 1 integers, to base 10 representation
         powers_of_two = 2 ** np.arange(len(device_wires))[::-1]
         indices = samples @ powers_of_two
 
         # count the basis state occurrences, and construct the probability vector
-        basis_states, counts = np.unique(indices, return_counts=True)
-        prob = np.zeros([2 ** len(device_wires)], dtype=np.float64)
-        prob[basis_states] = counts / len(samples)
+        if bin_size is not None:
+            bins = len(samples) // bin_size
+
+            indices = indices.reshape((bins, -1))
+            prob = np.zeros([2 ** len(device_wires), bins], dtype=np.float64)
+
+            # count the basis state occurrences, and construct the probability vector
+            for b, idx in enumerate(indices):
+                basis_states, counts = np.unique(idx, return_counts=True)
+                prob[basis_states, b] = counts / bin_size
+
+        else:
+            basis_states, counts = np.unique(indices, return_counts=True)
+            prob = np.zeros([2 ** len(device_wires)], dtype=np.float64)
+            prob[basis_states] = counts / len(samples)
+
         return self._asarray(prob, dtype=self.R_DTYPE)
 
-    def probability(self, wires=None):
+    def probability(self, wires=None, shot_range=None, bin_size=None):
         """Return either the analytic probability or estimated probability of
         each computational basis state.
 
@@ -591,7 +666,7 @@ class QubitDevice(Device):
         if hasattr(self, "analytic") and self.analytic:
             return self.analytic_probability(wires=wires)
 
-        return self.estimate_probability(wires=wires)
+        return self.estimate_probability(wires=wires, shot_range=shot_range, bin_size=bin_size)
 
     def marginal_prob(self, prob, wires=None):
         r"""Return the marginal probability of the computational basis
@@ -662,7 +737,7 @@ class QubitDevice(Device):
         perm = basis_states @ powers_of_two
         return self._gather(prob, perm)
 
-    def expval(self, observable):
+    def expval(self, observable, shot_range=None, bin_size=None):
 
         if self.analytic:
             # exact expectation value
@@ -671,9 +746,10 @@ class QubitDevice(Device):
             return self._dot(eigvals, prob)
 
         # estimate the ev
-        return np.mean(self.sample(observable))
+        samples = self.sample(observable, shot_range=shot_range, bin_size=bin_size)
+        return np.squeeze(np.mean(samples, axis=0))
 
-    def var(self, observable):
+    def var(self, observable, shot_range=None, bin_size=None):
 
         if self.analytic:
             # exact variance value
@@ -682,24 +758,34 @@ class QubitDevice(Device):
             return self._dot((eigvals ** 2), prob) - self._dot(eigvals, prob) ** 2
 
         # estimate the variance
-        return np.var(self.sample(observable))
+        samples = self.sample(observable, shot_range=shot_range, bin_size=bin_size)
+        return np.squeeze(np.var(samples, axis=0))
 
-    def sample(self, observable):
+    def sample(self, observable, shot_range=None, bin_size=None):
 
         # translate to wire labels used by device
         device_wires = self.map_wires(observable.wires)
         name = observable.name
+        sample_slice = Ellipsis if shot_range is None else slice(*shot_range)
 
         if isinstance(name, str) and name in {"PauliX", "PauliY", "PauliZ", "Hadamard"}:
             # Process samples for observables with eigenvalues {1, -1}
-            return 1 - 2 * self._samples[:, device_wires[0]]
+            samples = 1 - 2 * self._samples[sample_slice, device_wires[0]]
 
-        # Replace the basis state in the computational basis with the correct eigenvalue.
-        # Extract only the columns of the basis samples required based on ``wires``.
-        samples = self._samples[:, np.array(device_wires)]  # Add np.array here for Jax support.
-        powers_of_two = 2 ** np.arange(samples.shape[-1])[::-1]
-        indices = samples @ powers_of_two
-        return observable.eigvals[indices]
+        else:
+            # Replace the basis state in the computational basis with the correct eigenvalue.
+            # Extract only the columns of the basis samples required based on ``wires``.
+            samples = self._samples[
+                sample_slice, np.array(device_wires)
+            ]  # Add np.array here for Jax support.
+            powers_of_two = 2 ** np.arange(samples.shape[-1])[::-1]
+            indices = samples @ powers_of_two
+            samples = observable.eigvals[indices]
+
+        if bin_size is None:
+            return samples
+
+        return samples.reshape((bin_size, -1))
 
     def adjoint_jacobian(self, tape):
         """Implements the adjoint method outlined in


Test patch:
diff --git a/tests/test_qubit_device.py b/tests/test_qubit_device.py
index 03c0e4dd0..9cbc68635 100644
--- a/tests/test_qubit_device.py
+++ b/tests/test_qubit_device.py
@@ -45,10 +45,10 @@ def mock_qubit_device(monkeypatch):
         m.setattr(QubitDevice, "operations", ["PauliY", "RX", "Rot"])
         m.setattr(QubitDevice, "observables", ["PauliZ"])
         m.setattr(QubitDevice, "short_name", "MockDevice")
-        m.setattr(QubitDevice, "expval", lambda self, x: 0)
-        m.setattr(QubitDevice, "var", lambda self, x: 0)
-        m.setattr(QubitDevice, "sample", lambda self, x: 0)
-        m.setattr(QubitDevice, "apply", lambda self, x: None)
+        m.setattr(QubitDevice, "expval", lambda self, *args, **kwargs: 0)
+        m.setattr(QubitDevice, "var", lambda self, *args, **kwargs: 0)
+        m.setattr(QubitDevice, "sample", lambda self, *args, **kwargs: 0)
+        m.setattr(QubitDevice, "apply", lambda self, *args, **kwargs: None)
 
         def get_qubit_device(wires=1):
             return QubitDevice(wires=wires)
@@ -66,13 +66,13 @@ def mock_qubit_device_extract_stats(monkeypatch):
         m.setattr(QubitDevice, "operations", ["PauliY", "RX", "Rot"])
         m.setattr(QubitDevice, "observables", ["PauliZ"])
         m.setattr(QubitDevice, "short_name", "MockDevice")
-        m.setattr(QubitDevice, "expval", lambda self, x: 0)
-        m.setattr(QubitDevice, "var", lambda self, x: 0)
-        m.setattr(QubitDevice, "sample", lambda self, x: 0)
+        m.setattr(QubitDevice, "expval", lambda self, *args, **kwargs: 0)
+        m.setattr(QubitDevice, "var", lambda self, *args, **kwargs: 0)
+        m.setattr(QubitDevice, "sample", lambda self, *args, **kwargs: 0)
         m.setattr(QubitDevice, "state", 0)
         m.setattr(QubitDevice, "density_matrix", lambda self, wires=None: 0)
         m.setattr(
-            QubitDevice, "probability", lambda self, wires=None: 0
+            QubitDevice, "probability", lambda self, wires=None, *args, **kwargs: 0
         )
         m.setattr(QubitDevice, "apply", lambda self, x: x)
 
@@ -115,9 +115,9 @@ def mock_qubit_device_with_paulis_and_methods(monkeypatch):
         m.setattr(QubitDevice, "operations", mock_qubit_device_paulis)
         m.setattr(QubitDevice, "observables", mock_qubit_device_paulis)
         m.setattr(QubitDevice, "short_name", "MockDevice")
-        m.setattr(QubitDevice, "expval", lambda self, x: 0)
-        m.setattr(QubitDevice, "var", lambda self, x: 0)
-        m.setattr(QubitDevice, "sample", lambda self, x: 0)
+        m.setattr(QubitDevice, "expval", lambda self, *args, **kwargs: 0)
+        m.setattr(QubitDevice, "var", lambda self, *args, **kwargs: 0)
+        m.setattr(QubitDevice, "sample", lambda self, *args, **kwargs: 0)
         m.setattr(QubitDevice, "apply", lambda self, x, rotations: None)
 
         def get_qubit_device(wires=1):
@@ -135,9 +135,9 @@ def mock_qubit_device_with_paulis_rotations_and_methods(monkeypatch):
         m.setattr(QubitDevice, "operations", mock_qubit_device_paulis + mock_qubit_device_rotations)
         m.setattr(QubitDevice, "observables", mock_qubit_device_paulis)
         m.setattr(QubitDevice, "short_name", "MockDevice")
-        m.setattr(QubitDevice, "expval", lambda self, x: 0)
-        m.setattr(QubitDevice, "var", lambda self, x: 0)
-        m.setattr(QubitDevice, "sample", lambda self, x: 0)
+        m.setattr(QubitDevice, "expval", lambda self, *args, **kwargs: 0)
+        m.setattr(QubitDevice, "var", lambda self, *args, **kwargs: 0)
+        m.setattr(QubitDevice, "sample", lambda self, *args, **kwargs: 0)
         m.setattr(QubitDevice, "apply", lambda self, x: None)
 
         def get_qubit_device(wires=1):
@@ -500,8 +500,8 @@ class TestExpval:
 
         call_history = []
         with monkeypatch.context() as m:
-            m.setattr(QubitDevice, "sample", lambda self, obs: obs)
-            m.setattr("numpy.mean", lambda obs: obs)
+            m.setattr(QubitDevice, "sample", lambda self, obs, *args, **kwargs: obs)
+            m.setattr("numpy.mean", lambda obs, axis=None: obs)
             res = dev.expval(obs)
 
         assert res == obs
@@ -548,8 +548,8 @@ class TestVar:
 
         call_history = []
         with monkeypatch.context() as m:
-            m.setattr(QubitDevice, "sample", lambda self, obs: obs)
-            m.setattr("numpy.var", lambda obs: obs)
+            m.setattr(QubitDevice, "sample", lambda self, obs, *args, **kwargs: obs)
+            m.setattr("numpy.var", lambda obs, axis=None: obs)
             res = dev.var(obs)
 
         assert res == obs
@@ -861,3 +861,124 @@ class TestBatchExecution:
 
         assert len(res) == 3
         assert np.allclose(res[0], dev.execute(empty_tape), rtol=tol, atol=0)
+
+
+class TestShotList:
+    """Tests for passing shots as a list"""
+
+    shot_data = [
+        [[1, 2, 3, 10], [(1, 1), (2, 1), (3, 1), (10, 1)], (4,), 16],
+        [[1, 2, 2, 2, 10, 1, 1, 5, 1, 1, 1], [(1, 1), (2, 3), (10, 1), (1, 2), (5, 1), (1, 3)], (11,), 27],
+        [[10, 10, 10], [(10, 3)], (3,), 30],
+        [[(10, 3)], [(10, 3)], (3,), 30],
+    ]
+
+    @pytest.mark.parametrize("shot_list,shot_vector,expected_shape,total_shots", shot_data)
+    def test_single_expval(self, shot_list, shot_vector, expected_shape, total_shots):
+        """Test a single expectation value"""
+        dev = qml.device("default.qubit", wires=2, analytic=False, shots=shot_list)
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RX(x, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        res = circuit(0.5)
+
+        assert res.shape == expected_shape
+        assert circuit.device._shot_vector == shot_vector
+        assert circuit.device.shots == total_shots
+
+    shot_data = [
+        [[1, 2, 3, 10], [(1, 1), (2, 1), (3, 1), (10, 1)], (4, 2), 16],
+        [[1, 2, 2, 2, 10, 1, 1, 5, 1, 1, 1], [(1, 1), (2, 3), (10, 1), (1, 2), (5, 1), (1, 3)], (11, 2), 27],
+        [[10, 10, 10], [(10, 3)], (3, 2), 30],
+        [[(10, 3)], [(10, 3)], (3, 2), 30],
+    ]
+
+    @pytest.mark.parametrize("shot_list,shot_vector,expected_shape,total_shots", shot_data)
+    def test_multiple_expval(self, shot_list, shot_vector, expected_shape, total_shots):
+        """Test multiple expectation values"""
+        dev = qml.device("default.qubit", wires=2, analytic=False, shots=shot_list)
+
+        @qml.qnode(dev)
+        def circuit(x, y):
+            qml.RX(x, wires=0)
+            qml.RY(y, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0))
+
+        res = circuit(0.5, 0.1)
+
+        assert res.shape == expected_shape
+        assert circuit.device._shot_vector == shot_vector
+        assert circuit.device.shots == total_shots
+
+        # test gradient works
+        res = qml.jacobian(circuit)(0.5, 0.1)
+        assert res.shape == (2,) + expected_shape
+
+    shot_data = [
+        [[1, 2, 3, 10], [(1, 1), (2, 1), (3, 1), (10, 1)], (4, 4), 16],
+        [[1, 2, 2, 2, 10, 1, 1, 5, 1, 1, 1], [(1, 1), (2, 3), (10, 1), (1, 2), (5, 1), (1, 3)], (11, 4), 27],
+        [[10, 10, 10], [(10, 3)], (3, 4), 30],
+        [[(10, 3)], [(10, 3)], (3, 4), 30],
+    ]
+
+    @pytest.mark.parametrize("shot_list,shot_vector,expected_shape,total_shots", shot_data)
+    def test_probs(self, shot_list, shot_vector, expected_shape, total_shots):
+        """Test a probability return"""
+        dev = qml.device("default.qubit", wires=2, analytic=False, shots=shot_list)
+
+        @qml.qnode(dev)
+        def circuit(x, y):
+            qml.RX(x, wires=0)
+            qml.RY(y, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.probs(wires=[0, 1])
+
+        res = circuit(0.5, 0.1)
+
+        assert res.shape == expected_shape
+        assert circuit.device._shot_vector == shot_vector
+        assert circuit.device.shots == total_shots
+
+        # test gradient works
+        res = qml.jacobian(circuit)(0.5, 0.1)
+
+    shot_data = [
+        [[1, 2, 3, 10], [(1, 1), (2, 1), (3, 1), (10, 1)], (4, 2, 2), 16],
+        [[1, 2, 2, 2, 10, 1, 1, 5, 1, 1, 1], [(1, 1), (2, 3), (10, 1), (1, 2), (5, 1), (1, 3)], (11, 2, 2), 27],
+        [[10, 10, 10], [(10, 3)], (3, 2, 2), 30],
+        [[(10, 3)], [(10, 3)], (3, 2, 2), 30],
+    ]
+
+    @pytest.mark.parametrize("shot_list,shot_vector,expected_shape,total_shots", shot_data)
+    def test_multiple_probs(self, shot_list, shot_vector, expected_shape, total_shots):
+        """Test multiple probability returns"""
+        dev = qml.device("default.qubit", wires=2, analytic=False, shots=shot_list)
+
+        @qml.qnode(dev)
+        def circuit(x, y):
+            qml.RX(x, wires=0)
+            qml.RY(y, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.probs(wires=0), qml.probs(wires=1)
+
+        res = circuit(0.5, 0.1)
+
+        assert res.shape == expected_shape
+        assert circuit.device._shot_vector == shot_vector
+        assert circuit.device.shots == total_shots
+
+        # test gradient works
+        res = qml.jacobian(circuit)(0.5, 0.1)
+
+    def test_invalid_shot_list(self):
+        """Test exception raised if the shot list is the wrong type"""
+        with pytest.raises(qml.DeviceError, match="Shots must be"):
+            qml.device("default.qubit", wires=2, analytic=False, shots=0.5)
+
+        with pytest.raises(ValueError, match="Unknown shot sequence"):
+            qml.device("default.qubit", wires=2, analytic=False, shots=["a", "b", "c"])