Repository: PennyLaneAI/pennylane
Pull request #5544 (2024-04-18 17:00:11)

Cast samples to `int64` instead of `int8` in `qml.counts`
**Context:**
[sc-61314] Issue #5513 was being caused by samples that were converted to decimals to be cast to `int8`. This caused overflow with 8 or more wires, resulting in negative values being present in the counts dictionary.

**Description of the Change:**
* Update `qml.counts.process_samples` and `QubitDevice._samples_to_counts` to cast to `int64` instead of `int8`.

**Benefits:**
No more overflow issues with `qml.counts`.

**Possible Drawbacks:**

**Related GitHub Issues:**
#5513 

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5544?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.66%. Comparing base [(`a47d9bc`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/a47d9bc8298e7ed8e4813b4707299d8d327c6a1b?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`d7ffc8d`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5544?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 3 commits behind head on master.


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #5544      +/-   ##
==========================================
- Coverage   99.67%   99.66%   -0.01%     
==========================================
  Files         406      406              
  Lines       37881    37635     -246     
==========================================
- Hits        37758    37510     -248     
- Misses        123      125       +2     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5544?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-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 10af27fa5..3a98cce18 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -377,6 +377,9 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* `qml.counts` no longer returns negative samples when measuring 8 or more wires.
+  [(#5544)](https://github.com/PennyLaneAI/pennylane/pull/5544)
+
 * The `dynamic_one_shot` transform now works with broadcasting.
   [(#5473)](https://github.com/PennyLaneAI/pennylane/pull/5473)
 
diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index e47b6ef14..3044fafc8 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -1459,7 +1459,7 @@ class QubitDevice(Device):
         if mp.obs is None and not isinstance(mp.mv, MeasurementValue):
             # convert samples and outcomes (if using) from arrays to str for dict keys
             samples = np.array([sample for sample in samples if not np.any(np.isnan(sample))])
-            samples = qml.math.cast_like(samples, qml.math.int8(0))
+            samples = qml.math.cast_like(samples, qml.math.int64(0))
             samples = np.apply_along_axis(_sample_to_str, -1, samples)
             batched_ndims = 3  # no observable was provided, batched samples will have shape (batch_size, shots, len(wires))
             if mp.all_outcomes:
diff --git a/pennylane/measurements/counts.py b/pennylane/measurements/counts.py
index 341d4cd2a..f85396651 100644
--- a/pennylane/measurements/counts.py
+++ b/pennylane/measurements/counts.py
@@ -314,7 +314,7 @@ class CountsMP(SampleMeasurement):
             exp2 = 2 ** np.arange(num_wires - 1, -1, -1)
             samples = np.einsum("...i,i", samples, exp2)
             new_shape = samples.shape
-            samples = qml.math.cast_like(samples, qml.math.int8(0))
+            samples = qml.math.cast_like(samples, qml.math.int64(0))
             samples = list(map(convert, samples.ravel()))
             samples = np.array(samples).reshape(new_shape)
 


Test patch:
diff --git a/tests/measurements/test_counts.py b/tests/measurements/test_counts.py
index 550c65e05..085bccd3c 100644
--- a/tests/measurements/test_counts.py
+++ b/tests/measurements/test_counts.py
@@ -145,6 +145,20 @@ class TestProcessSamples:
         total_counts = sum(count for count in result.values())
         assert total_counts == 997
 
+    @pytest.mark.parametrize("n_wires", [5, 8, 10])
+    @pytest.mark.parametrize("all_outcomes", [True, False])
+    def test_counts_multi_wires_no_overflow(self, n_wires, all_outcomes):
+        """Test that binary strings for wire samples are not negative due to overflow."""
+        shots = 1000
+        total_wires = 10
+        samples = np.random.choice([0, 1], size=(shots, total_wires)).astype(np.float64)
+        result = qml.counts(wires=list(range(n_wires)), all_outcomes=all_outcomes).process_samples(
+            samples, wire_order=list(range(total_wires))
+        )
+
+        assert sum(result.values()) == shots
+        assert all(0 <= int(sample, 2) <= 2**n_wires for sample in result.keys())
+
     def test_counts_obs(self):
         """Test that the counts function outputs counts of the right size for observables"""
         shots = 1000
diff --git a/tests/test_qubit_device.py b/tests/test_qubit_device.py
index 93bfecc25..ca7ff3b44 100644
--- a/tests/test_qubit_device.py
+++ b/tests/test_qubit_device.py
@@ -1476,27 +1476,54 @@ class TestResourcesTracker:
         ]
 
 
-def test_samples_to_counts_with_nan():
-    """Test that the counts function disregards failed measurements (samples including
-    NaN values) when totalling counts"""
-    # generate 1000 samples for 2 wires, randomly distributed between 0 and 1
-    device = qml.device("default.qubit.legacy", wires=2, shots=1000)
-    device._state = [0.5 + 0.0j, 0.5 + 0.0j, 0.5 + 0.0j, 0.5 + 0.0j]
-    device._samples = device.generate_samples()
-    samples = device.sample(qml.measurements.CountsMP())
-
-    # imitate hardware return with NaNs (requires dtype float)
-    samples = qml.math.cast_like(samples, np.array([1.2]))
-    samples[0][0] = np.NaN
-    samples[17][1] = np.NaN
-    samples[850][0] = np.NaN
-
-    result = device._samples_to_counts(samples, mp=qml.measurements.CountsMP(), num_wires=2)
-
-    # no keys with NaNs
-    assert len(result) == 4
-    assert set(result.keys()) == {"00", "01", "10", "11"}
-
-    # # NaNs were not converted into "0", but were excluded from the counts
-    total_counts = sum(count for count in result.values())
-    assert total_counts == 997
+class TestSamplesToCounts:
+    """Tests for correctness of QubitDevice._samples_to_counts"""
+
+    def test_samples_to_counts_with_nan(self):
+        """Test that the counts function disregards failed measurements (samples including
+        NaN values) when totalling counts"""
+        # generate 1000 samples for 2 wires, randomly distributed between 0 and 1
+        device = qml.device("default.qubit.legacy", wires=2, shots=1000)
+        device._state = [0.5 + 0.0j, 0.5 + 0.0j, 0.5 + 0.0j, 0.5 + 0.0j]
+        device._samples = device.generate_samples()
+        samples = device.sample(qml.measurements.CountsMP())
+
+        # imitate hardware return with NaNs (requires dtype float)
+        samples = qml.math.cast_like(samples, np.array([1.2]))
+        samples[0][0] = np.NaN
+        samples[17][1] = np.NaN
+        samples[850][0] = np.NaN
+
+        result = device._samples_to_counts(samples, mp=qml.measurements.CountsMP(), num_wires=2)
+
+        # no keys with NaNs
+        assert len(result) == 4
+        assert set(result.keys()) == {"00", "01", "10", "11"}
+
+        # # NaNs were not converted into "0", but were excluded from the counts
+        total_counts = sum(result.values())
+        assert total_counts == 997
+
+    @pytest.mark.parametrize("all_outcomes", [True, False])
+    def test_samples_to_counts_with_many_wires(self, all_outcomes):
+        """Test that the counts function correctly converts wire samples to strings when
+        the number of wires is 8 or more."""
+        # generate 1000 samples for 10 wires, randomly distributed between 0 and 1
+        n_wires = 10
+        shots = 100
+        device = qml.device("default.qubit.legacy", wires=n_wires, shots=shots)
+        state = np.random.rand(*([2] * n_wires))
+        device._state = state / np.linalg.norm(state)
+        device._samples = device.generate_samples()
+        samples = device.sample(qml.measurements.CountsMP(all_outcomes=all_outcomes))
+
+        result = device._samples_to_counts(
+            samples, mp=qml.measurements.CountsMP(), num_wires=n_wires
+        )
+
+        # Check that keys are correct binary strings
+        assert all(0 <= int(sample, 2) <= 2**n_wires for sample in result.keys())
+
+        # # NaNs were not converted into "0", but were excluded from the counts
+        total_counts = sum(result.values())
+        assert total_counts == shots