Repository: PennyLaneAI/pennylane
Pull request #1232 (2021-04-17 22:05:47)

Fix using a shot_list with multiple shots for sampling
**Context:**

The introduction of partitioned shots allows multiple jobs to be expressed with nearly the same syntax. This allows defining the `shots` attribute of a device as a list of integers.

The circuit runs are separated for each component in the shot vector. The results of these circuit runs are then merged into a single array.

When using `qml.sample`, there can be cases when the shot vector has non-identical elements:
```python
import pennylane as qml

shots_list = [5, 10, 1000]
dev = qml.device("default.qubit", wires=2, shots=shots_list)

@qml.qnode(dev)
def circuit(a):
    qml.RX(a, wires=0)
    return qml.sample(qml.PauliZ(wires=0))

circuit(0.8)
```
Eventually, as per the original logic of using a shot vector, we would like to "merge" (more so use `np.stack`) the result of each circuit run. This is, however, erroneous on the above example:
```python
~/xanadu/pennylane/pennylane/math/numpy_box.py in stack(values, axis)
    104     @wrap_output
    105     def stack(values, axis=0):
--> 106         return np.stack(NumpyBox.unbox_list(values), axis=axis)
    107 
    108     @staticmethod

<__array_function__ internals> in stack(*args, **kwargs)

~/anaconda3/lib/python3.8/site-packages/numpy/core/shape_base.py in stack(arrays, axis, out)
    425     shapes = {arr.shape for arr in arrays}
    426     if len(shapes) != 1:
--> 427         raise ValueError('all input arrays must have the same shape')
    428 
    429     result_ndim = arrays[0].ndim + 1

ValueError: all input arrays must have the same shape
```

**Description of the Change:**
Adds conditions such that the results of sampling with a shot vector are not transformed into an array.

Using `qml.sample` with a shot vector returns a tuple consisting of the results obtained for each component of the shot vector (similar to when having a `return` statement with multiple components).

**Benefits:**
Using `qml.sample` once or multiple times with a shot vector doesn't raise an error.

**Possible Drawbacks:**
N/A

**Related GitHub Issues:**
N/A

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.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1232?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> :exclamation: No coverage uploaded for pull request base (`master@a3def61`). [Click here to learn what that means](https://docs.codecov.io/docs/error-reference?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#section-missing-base-commit).
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1232/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1232?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@            Coverage Diff            @@
##             master    #1232   +/-   ##
=========================================
  Coverage          ?   98.12%           
=========================================
  Files             ?      145           
  Lines             ?    10991           
  Branches          ?        0           
=========================================
  Hits              ?    10785           
  Misses            ?      206           
  Partials          ?        0           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1232?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/\_device.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1232/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19kZXZpY2UucHk=) | `96.70% <100.00%> (ø)` | |
| [pennylane/\_qubit\_device.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1232/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19xdWJpdF9kZXZpY2UucHk=) | `98.50% <100.00%> (ø)` | |
| [pennylane/qnode.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1232/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3Fub2RlLnB5) | `98.67% <100.00%> (ø)` | |

------

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

josh146: Thanks for catching this @antalszava!

Quick question, could this be fixed in less lines by adding the following `if` statement?

```python
# compute the required statistics
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

    if not circuit.is_sampled:  # add this if statement here
        results = qml.math.stack(results)
```

With this `if` statement, the stacking is avoided if the circuit is being sampled, and the results are returned as a (ragged) list.
antalszava: @josh146, thanks for the suggestion! There were two reasons why it seemed less preferable:
* Creating the output emits the `VisibleDeprecationWarning` as we saw in [Strawberry Fields recently when creating ragged numpy arrays](https://github.com/XanaduAI/strawberryfields/pull/564) (we have to specify `'dtype=object'` explicitly to silence the warning)
* Personally, it seems more intuitive to me for a user to receive a tuple of the results, rather than a ragged array (whose `dtype` is also not describing its true content). However, if there are any advantages to having a ragged array, then happy to go with that approach.
josh146: > Personally, it seems more intuitive to me for a user to receive a tuple of the results, rather than a ragged array (whose dtype is also not describing its true content).

Yes, I agree a tuple is more intuitive - I didn't realize my suggestion _wasn't_ returning a tuple 😆 In that case, I am happy with your solution @antalszava. Although, I might make `has_partitioned_shots` private, not public?

Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index 977add323..45124b681 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -483,6 +483,9 @@
   transforms.
   [(#1064)](https://github.com/PennyLaneAI/pennylane/pull/1064)
 
+* Sampling works correctly with batches of shots specified as a list.
+  [(#1232)](https://github.com/PennyLaneAI/pennylane/pull/1232)
+
 * Adds a new transform `qml.invisible`.
   [(#1175)](https://github.com/PennyLaneAI/pennylane/pull/1175)
 
diff --git a/pennylane/_device.py b/pennylane/_device.py
index 631a37465..737dbb8b6 100644
--- a/pennylane/_device.py
+++ b/pennylane/_device.py
@@ -296,6 +296,16 @@ class Device(abc.ABC):
         """
         return self._shot_vector
 
+    def _has_partitioned_shots(self):
+        """Checks if the device was instructed to perform executions with partitioned shots.
+
+        Returns:
+            bool: whether or not shots are partitioned
+        """
+        return self._shot_vector is not None and (
+            len(self._shot_vector) > 1 or self._shot_vector[0].copies > 1
+        )
+
     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 9ae321caa..93d38a07e 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -187,6 +187,8 @@ class QubitDevice(Device):
         if self.shots is not None or circuit.is_sampled:
             self._samples = self.generate_samples()
 
+        multiple_sampled_jobs = circuit.is_sampled and self._has_partitioned_shots()
+
         # compute the required statistics
         if not self.analytic and self._shot_vector is not None:
 
@@ -203,16 +205,18 @@ class QubitDevice(Device):
                 if shot_tuple.copies > 1:
                     results.extend(r.T)
                 else:
-                    results.append(r)
+                    results.append(r.T)
 
                 s1 = s2
 
-            results = qml.math.stack(results)
+            if not multiple_sampled_jobs:
+                # Can only stack single element outputs
+                results = qml.math.stack(results)
 
         else:
             results = self.statistics(circuit.observables)
 
-        if circuit.all_sampled or not circuit.is_sampled:
+        if (circuit.all_sampled or not circuit.is_sampled) and not multiple_sampled_jobs:
             results = self._asarray(results)
         else:
             results = tuple(self._asarray(r) for r in results)
diff --git a/pennylane/qnode.py b/pennylane/qnode.py
index 3c00fd86d..abf4e5505 100644
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -573,7 +573,9 @@ class QNode:
             if hasattr(self._original_device, "_state"):
                 self._original_device._state = self.device._state
 
-        if isinstance(self.qfunc_output, Sequence):
+        if isinstance(self.qfunc_output, Sequence) or (
+            self.qtape.is_sampled and self.device._has_partitioned_shots()
+        ):
             return res
 
         return qml.math.squeeze(res)


Test patch:
diff --git a/tests/test_qubit_device.py b/tests/test_qubit_device.py
index d104c8cf2..94ec3b94a 100644
--- a/tests/test_qubit_device.py
+++ b/tests/test_qubit_device.py
@@ -930,6 +930,62 @@ class TestShotList:
         # 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)], [(), (2,), (3,), (10,)], 16],
+        [[1, 2, 2, 2, 10, 1, 1, 5, 1, 1, 1], [(1, 1), (2, 3), (10, 1), (1, 2), (5, 1), (1, 3)],
+                [(), (2,), (2,),(2,),(10,), (), (), (5,), (), (), ()], 27],
+        [[10, 10, 10], [(10, 3)], [(10,), (10,), (10,)], 30],
+        [[(10, 3)], [(10, 3)], [(10,),(10,),(10,)], 30],
+    ]
+
+    @pytest.mark.parametrize("shot_list,shot_vector,expected_shapes,total_shots", shot_data)
+    def test_sample(self, shot_list, shot_vector, expected_shapes, total_shots):
+        """Test sample returns"""
+        dev = qml.device("default.qubit", wires=2, 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.sample(qml.PauliZ(wires=0))
+
+        res = circuit(0.5, 0.1)
+
+        for r, shape in zip(res, expected_shapes):
+            assert r.shape == 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)], [(2,), (2,2), (3,2), (10,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)],
+                [(2,), (2,2), (2,2),(2,2),(10,2), (2,), (2,), (5,2), (2,), (2,), (2,)], 27],
+        [[10, 10, 10], [(10,3)], [(10,2), (10,2), (10,2)], 30],
+        [[(10, 3)], [(10, 3)], [(10,2),(10,2),(10,2)], 30],
+    ]
+
+    @pytest.mark.parametrize("shot_list,shot_vector,expected_shapes,total_shots", shot_data)
+    def test_multiple_sample(self, shot_list, shot_vector, expected_shapes, total_shots):
+        """Test sample returns"""
+        dev = qml.device("default.qubit", wires=2, 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.sample(qml.PauliZ(wires=0)), qml.sample(qml.PauliZ(wires=1))
+
+        res = circuit(0.5, 0.1)
+
+        for r, shape in zip(res, expected_shapes):
+            assert r.shape == shape
+
+        assert circuit.device._shot_vector == shot_vector
+        assert circuit.device.shots == total_shots
+
     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"):