Repository: PennyLaneAI/pennylane
Pull request #5556 (2024-04-22 21:11:41)

Fix `qml.counts` overflow error for more than 64 qubits
**Context:**
`qml.counts` was converting samples to integers and using f-strings to format as binary strings, however, it caused wires > 64 to have overflow.

Without the changes in this PR:
```python
import pennylane as qml
import numpy as np

np.random.seed(582)

num_wires = 70
wires = list(range(num_wires))

dev = qml.device("default.clifford", wires=wires, shots=100)

@qml.qnode(dev)
# @qml.defer_measurements
def circuit():
    [qml.PauliX(w) for w in wires]
    return qml.counts()

circuit()
```
```pycon
{'-000000000000000000000000000000000000000000000000000000000000000000001': tensor(100, requires_grad=True)}
```

With the changes:
```python
import pennylane as qml
import numpy as np

np.random.seed(582)

num_wires = 70
wires = list(range(num_wires))

dev = qml.device("default.clifford", wires=wires, shots=100)

@qml.qnode(dev)
# @qml.defer_measurements
def circuit():
    [qml.PauliX(w) for w in wires]
    return qml.counts()

circuit()
```
```pycon
{'1111111111111111111111111111111111111111111111111111111111111111111111': tensor(100, requires_grad=True)}
```

**Description of the Change:**
* Change `qml.counts` to convert directly from list of samples to binary string without first converting to integers.
* Changed `QubitDevice._samples_to_counts` to cast to `int8` instead of `int64`. This is only for casting samples, not converting to ints, as `int8` is more memory efficient.

**Benefits:**
* `qml.counts` now works with arbitrary wires. This is useful for devices like `default.clifford`, which can simulate hundreds of qubits.

**Possible Drawbacks:**

**Related GitHub Issues:**
#5537 , #5513 

mudit2812: [sc-61314]
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5556?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.68%. Comparing base [(`a8bba50`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/a8bba509e16a8582d57a0baa978ddb89a6d4d75a?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`289ea82`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5556?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 1 commits behind head on master.


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


```diff
@@            Coverage Diff             @@
##           master    #5556      +/-   ##
==========================================
- Coverage   99.69%   99.68%   -0.01%     
==========================================
  Files         410      410              
  Lines       38230    37939     -291     
==========================================
- Hits        38113    37821     -292     
- Misses        117      118       +1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5556?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 563e74dd0..1289ec086 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -470,6 +470,7 @@
 
 * `qml.counts` no longer returns negative samples when measuring 8 or more wires.
   [(#5544)](https://github.com/PennyLaneAI/pennylane/pull/5544)
+  [(#5556)](https://github.com/PennyLaneAI/pennylane/pull/5556)
 
 * 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 ba8b03121..dbb01c0e0 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -1462,7 +1462,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.int64(0))
+            samples = qml.math.cast_like(samples, qml.math.int8(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 f85396651..7a99ae84c 100644
--- a/pennylane/measurements/counts.py
+++ b/pennylane/measurements/counts.py
@@ -299,6 +299,7 @@ class CountsMP(SampleMeasurement):
 
         if self.obs is None and not isinstance(self.mv, MeasurementValue):
             # convert samples and outcomes (if using) from arrays to str for dict keys
+            batched_ndims = 3  # no observable was provided, batched samples will have shape (batch_size, shots, len(wires))
 
             # remove nans
             mask = qml.math.isnan(samples)
@@ -307,21 +308,24 @@ class CountsMP(SampleMeasurement):
                 mask = np.logical_not(np.any(mask, axis=tuple(range(1, samples.ndim))))
                 samples = samples[mask, ...]
 
-            # convert to string
-            def convert(x):
-                return f"{x:0{num_wires}b}"
+            def convert(sample):
+                # convert array of ints to string
+                return "".join(str(s) for s in sample)
 
-            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.int64(0))
-            samples = list(map(convert, samples.ravel()))
-            samples = np.array(samples).reshape(new_shape)
+            new_shape = samples.shape[:-1]
+            # Flatten broadcasting axis
+            flattened_samples = np.reshape(samples, (-1, shape[-1])).astype(np.int8)
+            samples = list(map(convert, flattened_samples))
+            samples = np.reshape(np.array(samples), new_shape)
 
-            batched_ndims = 3  # no observable was provided, batched samples will have shape (batch_size, shots, len(wires))
             if self.all_outcomes:
+
+                def convert_from_int(x):
+                    # convert int to binary string
+                    return f"{x:0{num_wires}b}"
+
                 num_wires = len(self.wires) if len(self.wires) > 0 else shape[-1]
-                outcomes = list(map(convert, range(2**num_wires)))
+                outcomes = list(map(convert_from_int, range(2**num_wires)))
 
         elif self.all_outcomes:
             # This also covers statistics for mid-circuit measurements manipulated using
diff --git a/pennylane/transforms/dynamic_one_shot.py b/pennylane/transforms/dynamic_one_shot.py
index 98625aa60..2480bcff9 100644
--- a/pennylane/transforms/dynamic_one_shot.py
+++ b/pennylane/transforms/dynamic_one_shot.py
@@ -381,5 +381,5 @@ def gather_mcm(measurement, samples):
         meas_tmp = measurement.__class__(wires=wires)
     new_measurement = meas_tmp.process_samples(mcm_samples, wire_order=wires)
     if isinstance(measurement, CountsMP) and not use_as_is:
-        new_measurement = dict(sorted((int(x, 2), y) for x, y in new_measurement.items()))
+        new_measurement = dict(sorted((int(x), y) for x, y in new_measurement.items()))
     return new_measurement


Test patch:
diff --git a/tests/measurements/test_counts.py b/tests/measurements/test_counts.py
index 085bccd3c..afdba1ed7 100644
--- a/tests/measurements/test_counts.py
+++ b/tests/measurements/test_counts.py
@@ -145,19 +145,29 @@ 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):
+    @pytest.mark.parametrize("batch_size", [None, 1, 4])
+    @pytest.mark.parametrize("n_wires", [4, 10, 65])
+    @pytest.mark.parametrize("all_outcomes", [False, True])
+    def test_counts_multi_wires_no_overflow(self, n_wires, all_outcomes, batch_size):
         """Test that binary strings for wire samples are not negative due to overflow."""
+        if all_outcomes and n_wires == 65:
+            pytest.skip("Too much memory being used, skipping")
         shots = 1000
-        total_wires = 10
-        samples = np.random.choice([0, 1], size=(shots, total_wires)).astype(np.float64)
+        total_wires = 65
+        shape = (batch_size, shots, total_wires) if batch_size else (shots, total_wires)
+        samples = np.random.choice([0, 1], size=shape).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())
+        if batch_size:
+            assert len(result) == batch_size
+            for r in result:
+                assert sum(r.values()) == shots
+                assert all("-" not in sample for sample in r.keys())
+        else:
+            assert sum(result.values()) == shots
+            assert all("-" not in sample for sample in result.keys())
 
     def test_counts_obs(self):
         """Test that the counts function outputs counts of the right size for observables"""