Repository: PennyLaneAI/pennylane

Issue:
[BUG] `qml.counts` of a mid-circuit measurement does not return all outcomes even when requested
### Expected behavior

`qml.counts` with `all_outcomes=True` should return all outcomes.

### Actual behavior

The correct result is returned with deferred measurement.
```python
@qml.qnode(device=dev, mcm_method="deferred")
def circuit():
    m = qml.measure(0)
    return qml.counts(m, all_outcomes=True)
```
```pycon
>>> circuit()
{0: 10, 1: 0}
```
With `"one-shot"`, we get:
```python
@qml.qnode(device=dev, mcm_method="one-shot")
def circuit():
    m = qml.measure(0)
    return qml.counts(m, all_outcomes=True)
```
```pycon
>>> circuit()
{0.0: 10}
```
We see that:
1. We get a float `0.0` instead of the expected integer `0`.
2. We only get one outcome.

### Additional information

_No response_

### Source code

_No response_

### Tracebacks

_No response_

### System information

```shell
Name: PennyLane
Version: 0.40.0.dev31
Summary: PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Train a quantum computer the same way as a neural network.
Home-page: https://github.com/PennyLaneAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /Users/REDACTED/pennylane/venv/lib/python3.10/site-packages
Requires: appdirs, autograd, autoray, cachetools, diastatic-malt, networkx, numpy, packaging, pennylane-lightning, requests, rustworkx, scipy, tomlkit, typing_extensions
Required-by: PennyLane-Catalyst, PennyLane_Lightning, PennyLane_Lightning_Kokkos

Platform info:           macOS-15.1.1-arm64-arm-64bit
Python version:          3.10.14
Numpy version:           1.26.4
Scipy version:           1.12.0
Installed devices:
- default.clifford (PennyLane-0.40.0.dev32)
- default.gaussian (PennyLane-0.40.0.dev32)
- default.mixed (PennyLane-0.40.0.dev32)
- default.qubit (PennyLane-0.40.0.dev32)
- default.qutrit (PennyLane-0.40.0.dev32)
- default.qutrit.mixed (PennyLane-0.40.0.dev32)
- default.tensor (PennyLane-0.40.0.dev32)
- null.qubit (PennyLane-0.40.0.dev32)
- reference.qubit (PennyLane-0.40.0.dev32)
- nvidia.custatevec (PennyLane-Catalyst-0.9.0.dev41)
- nvidia.cutensornet (PennyLane-Catalyst-0.9.0.dev41)
- oqc.cloud (PennyLane-Catalyst-0.9.0.dev41)
- softwareq.qpp (PennyLane-Catalyst-0.9.0.dev41)
- lightning.kokkos (PennyLane_Lightning_Kokkos-0.39.0.dev38)
- lightning.qubit (PennyLane_Lightning-0.39.0)
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.

Discussion:

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 64f737f63..71ad0fe84 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -625,6 +625,9 @@ same information.
 
 <h3>Bug fixes 🐛</h3>
 
+* `qml.counts` returns all outcomes when the `all_outcomes` argument is `True` and mid-circuit measurements are present.
+  [(#6732)](https://github.com/PennyLaneAI/pennylane/pull/6732)
+
 * `qml.ControlledQubitUnitary` has consistent behaviour with program capture enabled. 
   [(#6719)](https://github.com/PennyLaneAI/pennylane/pull/6719)
 
diff --git a/pennylane/transforms/dynamic_one_shot.py b/pennylane/transforms/dynamic_one_shot.py
index a9fc151b7..4328c21bd 100644
--- a/pennylane/transforms/dynamic_one_shot.py
+++ b/pennylane/transforms/dynamic_one_shot.py
@@ -385,6 +385,17 @@ def gather_non_mcm(measurement, samples, is_valid, postselect_mode=None):
     """
     if isinstance(measurement, CountsMP):
         tmp = Counter()
+
+        if measurement.all_outcomes:
+            if isinstance(measurement.mv, Sequence):
+                values = [list(m.branches.values()) for m in measurement.mv]
+                values = list(itertools.product(*values))
+                tmp = Counter({"".join(map(str, v)): 0 for v in values})
+            else:
+                values = [list(measurement.mv.branches.values())]
+                values = list(itertools.product(*values))
+                tmp = Counter({float(*v): 0 for v in values})
+
         for i, d in enumerate(samples):
             tmp.update(
                 {k if isinstance(k, str) else float(k): v * is_valid[i] for k, v in d.items()}


Test patch:
diff --git a/tests/measurements/test_counts.py b/tests/measurements/test_counts.py
index e49ab0427..cae5816a3 100644
--- a/tests/measurements/test_counts.py
+++ b/tests/measurements/test_counts.py
@@ -396,6 +396,35 @@ class TestProcessSamples:
 class TestCountsIntegration:
     # pylint:disable=too-many-public-methods,not-an-iterable
 
+    def test_counts_all_outcomes_with_mcm(self):
+        """Test that all outcomes are present in results if requested."""
+        n_sample = 10
+
+        dev = qml.device("default.qubit", shots=n_sample)
+
+        @qml.qnode(device=dev, mcm_method="one-shot")
+        def single_mcm():
+            m = qml.measure(0)
+            return qml.counts(m, all_outcomes=True)
+
+        res = single_mcm()
+
+        assert list(res.keys()) == [0.0, 1.0]
+        assert sum(res.values()) == n_sample
+        assert res[0.0] == n_sample
+
+        @qml.qnode(device=dev, mcm_method="one-shot")
+        def double_mcm():
+            m1 = qml.measure(0)
+            m2 = qml.measure(1)
+            return qml.counts([m1, m2], all_outcomes=True)
+
+        res = double_mcm()
+
+        assert list(res.keys()) == ["00", "01", "10", "11"]
+        assert sum(res.values()) == n_sample
+        assert res["00"] == n_sample
+
     def test_counts_dimension(self):
         """Test that the counts function outputs counts of the right size"""
         n_sample = 10
diff --git a/tests/test_compiler.py b/tests/test_compiler.py
index 27a50a837..38d61cd86 100644
--- a/tests/test_compiler.py
+++ b/tests/test_compiler.py
@@ -864,8 +864,7 @@ class TestCatalystMCMs:
     @pytest.mark.parametrize(
         "measure_f",
         [
-            # https://github.com/PennyLaneAI/pennylane/issues/6700
-            pytest.param(qml.counts, marks=pytest.mark.xfail),
+            qml.counts,
             qml.expval,
             qml.probs,
         ],
@@ -898,7 +897,7 @@ class TestCatalystMCMs:
             meas_key = "wires" if isinstance(meas_obj, list) else "op"
             meas_value = m0 if isinstance(meas_obj, str) else meas_obj
             kwargs = {meas_key: meas_value}
-            if measure_f == qml.counts:
+            if measure_f is qml.counts:
                 kwargs["all_outcomes"] = True
             return measure_f(**kwargs)
 
@@ -919,17 +918,20 @@ class TestCatalystMCMs:
             meas_key = "wires" if isinstance(meas_obj, list) else "op"
             meas_value = m0 if isinstance(meas_obj, str) else meas_obj
             kwargs = {meas_key: meas_value}
-            if measure_f == qml.counts:
+            if measure_f is qml.counts:
                 kwargs["all_outcomes"] = True
             return measure_f(**kwargs)
 
         params = jnp.pi / 4 * jnp.ones(2)
         results0 = ref_func(*params)
         results1 = func(*params)
-        if measure_f == qml.counts:
-            ndim = 2  # both [0] and m0 are on one wire only
-            results1 = {format(int(state), f"0{ndim}b"): count for state, count in zip(*results1)}
-        if measure_f == qml.sample:
+        if measure_f is qml.counts:
+
+            def fname(x):
+                return format(x, f"0{len(meas_obj)}b") if isinstance(meas_obj, list) else x
+
+            results1 = {fname(int(state)): count for state, count in zip(*results1)}
+        if measure_f is qml.sample:
             results0 = results0[results0 != fill_in_value]
             results1 = results1[results1 != fill_in_value]
         mcm_utils.validate_measurements(measure_f, shots, results1, results0)