[SHOTS] Allows a list of shots to be specified

.github/CHANGELOG.md

@@ -2,6 +2,42 @@
 
 <h3>New features since last release</h3>
 
+* 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)
 

pennylane/_device.py

@@ -16,7 +16,7 @@
 """
 # pylint: disable=too-many-format-args
 import abc
-from collections.abc import Iterable
+from collections.abc import Iterable, Sequence
 from collections import OrderedDict
 
 import numpy as np
@@ -36,6 +36,48 @@
 from pennylane.wires import Wires, WireError
 
 
+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[int, tuple[int]]]: A tuple containing the total number
+        of shots, as well as the shot vector.
+
+    **Example**
+
+    >>> shot_list = [3, 1, 2, 2, 2, 2, 6, 1, 1, 5, 12, 10, 10]
+    >>> _process_shot_sequence(shot_list)
+    (57, [3, 1, (2, 4), 6, (1, 2), 5, 12, (10, 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; represent the shot vector
+            # in a sparse format.
+            shot_vector = [(shot_list[0], len(shot_list))]
+            total_shots = shot_list[0] * len(shot_list)
+        else:
+            # Iterate through the shots, and group consecutive identical shots
+            split_at_repeated = np.split(shot_list, np.where(np.diff(shot_list) != 0)[0] + 1)
+            shot_vector = [i.item() if len(i) == 1 else (i[0], len(i)) for i in split_at_repeated]
+            total_shots = np.sum(shot_list)
+
+    else:
+        # shot_list is already a shot vector, simply compute the total number of shots
+        shot_vector = shot_list
+        total_shots = sum(i if isinstance(i, int) else np.prod(i) for i in shot_list)
+
+    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 +220,23 @@ def shots(self, shots):
         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.
 

pennylane/_qubit_device.py

@@ -17,7 +17,7 @@
 
 # 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
@@ -208,7 +208,34 @@ def execute(self, circuit, **kwargs):
             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 in self._shot_vector:
+
+                if isinstance(shot, tuple):
+                    s2 = s1 + np.prod(shot)
+                    bin_size = shot[0]
+                else:
+                    s2 = s1 + shot
+                    bin_size = shot
+
+                r = self.statistics(circuit.observables, shot_range=[s1, s2], bin_size=bin_size)
+                r = np.squeeze(r)
+
+                if isinstance(shot, tuple):
+                    results.extend(r.T)
+                else:
+                    results.append(np.array(r))
+
+                s1 = s2
+
+            results = np.stack(results)
+
+        else:
+            results = self.statistics(circuit.observables)
 
         if circuit.all_sampled or not circuit.is_sampled:
             results = self._asarray(results)
@@ -317,14 +344,19 @@ def active_wires(operators):
 
         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
         density matrices.
 
         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
@@ -337,16 +369,18 @@ def statistics(self, observables):
         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 +459,7 @@ def sample_basis_states(self, number_of_states, state_probability):
 
         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 +578,18 @@ def analytic_probability(self, wires=None):
         """
         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 +601,33 @@ def estimate_probability(self, wires=None):
         # 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 +645,7 @@ def probability(self, wires=None):
         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 +716,7 @@ def marginal_prob(self, prob, wires=None):
         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 +725,10 @@ def expval(self, observable):
             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 +737,34 @@ def var(self, observable):
             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