qml.counts returns all outcomes when requested if MCMs present

doc/releases/changelog-dev.md

@@ -511,6 +511,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)
 

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()}

tests/measurements/test_counts.py

@@ -396,6 +396,34 @@ def test_counts_binsize(self):
 class TestCountsIntegration:
     # pylint:disable=too-many-public-methods,not-an-iterable
 
+    def test_counts_all_outcomes_with_mcm(self):
+        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

tests/test_compiler.py

@@ -481,7 +481,8 @@ def inner(j):
         x = 0.5
         assert jnp.allclose(circuit(x, 3), qml.qjit(circuit)(x, 3))
 
-        res = circuit.tape.operations
+        tape = qml.workflow.construct_tape(circuit)(x, 3)
+        res = tape.operations
         expected = [
             qml.Hadamard(wires=[0]),
             qml.Hadamard(wires=[1]),
@@ -864,8 +865,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,
         ],
@@ -927,8 +927,10 @@ def ansatz():
         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)}
+            # Convert keys to the same data-type for comparison
+            results1 = {format(int(state), "b"): count for state, count in zip(*results1)}
+            if all(not isinstance(key, str) for key in results0.keys()):
+                results0 = {format(state, "b"): count for state, count in results0.items()}
         if measure_f == qml.sample:
             results0 = results0[results0 != fill_in_value]
             results1 = results1[results1 != fill_in_value]