Repository: PennyLaneAI/pennylane (python)

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.

Pull request description:
`qml.counts` returns all outcomes when requested if `MCM`s present
**Context:**

Previously, despite using `all_outcomes` , `qml.counts` incorrectly returned only the states that have counts present,
```python
import pennylane as qml
dev = qml.device("default.qubit", shots=10)

@qml.qnode(device=dev, mcm_method="one-shot")
def circuit():
    m1 = qml.measure(0)
    return qml.counts(m1, all_outcomes=True)

>>> circuit()
{0.0: 10}
```

**Description of the Change:**

If `all_outcomes` is specified, the intermediate `tmp` variable (used to keep track of measurement values from `MCM`s) defaults to a `Counter` object that initializes all possible basis states as keys with their values set to zero. This counter is then updated based on the measurement results, where each mcm sample contributes to updating the counts, weighted by a validity factor.

This results in the correct output,
```python
dev = qml.device("default.qubit", shots=10)

@qml.qnode(device=dev, mcm_method="one-shot")
def circuit():
    m1 = qml.measure(0)
    return qml.counts(m1, all_outcomes=True)

>>> circuit()
{0.0: 10, 1.0: 0}
# {0.0: 10} if all_outcomes=False
```
or,
```python
dev = qml.device("default.qubit", shots=10)

@qml.qnode(device=dev, mcm_method="one-shot")
def circuit():
    m1 = qml.measure(0)
    m2 = qml.measure(1)
    return qml.counts([m1,m2], all_outcomes=True)

>>> circuit()
{'00': 10, '01': 0, '10': 0, '11': 0}
# {'00': 10} if all_outcomes=False
```

**Benefits:** `all_outcomes` kwarg in `qml.counts` now works as intended when `mcm_method="one-shot"` and `MCM` operations are present.

**Possible Drawbacks:** `tmp` object creation doesn't scale well with higher qubit counts.

**Related GitHub Issues:** Fixes #6700 

[sc-80206]

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)