Repository: PennyLaneAI/pennylane
Pull request #3422 (2022-11-25 20:48:34)

Fix `qml.sample` for shot vectors
**Context:**
Previously `qml.sample` produced incorrect results when shot vectors are used:
```python3
dev = qml.device("default.qubit", wires=2, shots=(5, 10))

@qml.qnode(dev)
def circuit():
    qml.Hadamard(wires=0)
    qml.CNOT(wires=[0, 1])
    return qml.sample(wires=[0, 1])
```
```
>>> circuit()
(tensor([[0, 0],
        [0, 1],
        [0, 1],
        [0, 1],
        [0, 1]], dtype=int64, requires_grad=True), tensor([[1, 0],
        [1, 0],
        [1, 0],
        [1, 0],
        [0, 0],
        [0, 0],
        [1, 0],
        [1, 0],
        [1, 1],
        [1, 1]], dtype=int64, requires_grad=True))
```

**Description of the Change:**
The reshape logic in the qubit device was missing a transpose. This PR adds the transpose. Also fixed a broken test that would have failed on master.

**Related GitHub Issues:**
https://github.com/PennyLaneAI/pennylane/issues/3412


github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.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/3422?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3422](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3422?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (995736e) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/fa0eb863e0bf87f7f05fe3aeacf9f093be507b11?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (fa0eb86) will **not change** coverage.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master    #3422   +/-   ##
=======================================
  Coverage   99.69%   99.69%           
=======================================
  Files         299      299           
  Lines       26617    26617           
=======================================
  Hits        26536    26536           
  Misses         81       81           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3422?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/\_qubit\_device.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3422/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19xdWJpdF9kZXZpY2UucHk=) | `99.66% <ø> (ø)` | |
| [pennylane/measurements/counts.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3422/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmVtZW50cy9jb3VudHMucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/measurements/sample.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3422/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmVtZW50cy9zYW1wbGUucHk=) | `100.00% <100.00%> (ø)` | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 0a0ad3fd8..c5670f455 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -371,6 +371,10 @@ Deprecations cycles are tracked at [doc/developement/deprecations.rst](https://d
   with all interfaces
   [(#3392)](https://github.com/PennyLaneAI/pennylane/pull/3392)
 
+* Fixed a bug where a QNode returning `qml.sample` would produce incorrect results when
+  run on a device defined with a shot vector.
+  [#3422](https://github.com/PennyLaneAI/pennylane/pull/3422)
+
 <h3>Contributors</h3>
 
 This release contains contributions from (in alphabetical order):
diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index 01f0dc6d3..0213845c1 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -1756,7 +1756,7 @@ class QubitDevice(Device):
             ]
 
         return (
-            samples.reshape((num_wires, bin_size, -1))
+            samples.T.reshape((num_wires, bin_size, -1))
             if no_observable_provided
             else samples.reshape((bin_size, -1))
         )
diff --git a/pennylane/measurements/counts.py b/pennylane/measurements/counts.py
index b423aa4c0..d725cceb0 100644
--- a/pennylane/measurements/counts.py
+++ b/pennylane/measurements/counts.py
@@ -178,9 +178,13 @@ class _Counts(SampleMeasurement):
             return self._samples_to_counts(samples)
 
         num_wires = len(self.wires) if self.wires else len(wire_order)
-        shape = (-1, bin_size, num_wires) if self.obs is None else (-1, bin_size)
+        samples = (
+            samples.reshape((num_wires, -1)).T.reshape(-1, bin_size, num_wires)
+            if self.obs is None
+            else samples.reshape((-1, bin_size))
+        )
 
-        return [self._samples_to_counts(bin_sample) for bin_sample in samples.reshape(shape)]
+        return [self._samples_to_counts(bin_sample) for bin_sample in samples]
 
     def _samples_to_counts(self, samples):
         """Groups the samples into a dictionary showing number of occurences for
diff --git a/pennylane/measurements/sample.py b/pennylane/measurements/sample.py
index 425f35f9b..fb930acfe 100644
--- a/pennylane/measurements/sample.py
+++ b/pennylane/measurements/sample.py
@@ -141,7 +141,7 @@ class _Sample(SampleMeasurement):
 
         if self.obs is None:
             # if no observable was provided then return the raw samples
-            return samples if bin_size is None else samples.reshape(num_wires, bin_size, -1)
+            return samples if bin_size is None else samples.T.reshape(num_wires, bin_size, -1)
 
         if name in {"PauliX", "PauliY", "PauliZ", "Hadamard"}:
             # Process samples for observables with eigenvalues {1, -1}


Test patch:
diff --git a/tests/measurements/test_sample.py b/tests/measurements/test_sample.py
index ed499a493..3adf8c5f5 100644
--- a/tests/measurements/test_sample.py
+++ b/tests/measurements/test_sample.py
@@ -230,15 +230,21 @@ class TestSample:
         @qml.qnode(dev)
         def circuit():
             qml.Hadamard(wires=0)
+            qml.CNOT(wires=[0, 1])
             return qml.sample()
 
         res = circuit()
 
         assert isinstance(res, tuple)
 
-        expected_shapes = [(num_wires,), (num_wires, shots2), (num_wires, shots3)]
+        expected_shapes = [(num_wires,), (shots2, num_wires), (shots3, num_wires)]
         assert len(res) == len(expected_shapes)
-        assert (r.shape == exp_shape for r, exp_shape in zip(res, expected_shapes))
+        assert all(r.shape == exp_shape for r, exp_shape in zip(res, expected_shapes))
+
+        # assert first wire is always the same as second
+        assert np.all(res[0][0] == res[0][1])
+        assert np.all(res[1][:, 0] == res[1][:, 1])
+        assert np.all(res[2][:, 0] == res[2][:, 1])
 
         custom_measurement_process(dev, spy)