Patch dynamic_one_shot for single shot case

doc/releases/changelog-0.38.0.md

@@ -375,6 +375,8 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* Fixes an error in the `dynamic_one_shot` transform when used with sampling a single shot.
+
 * `qml.transforms.broadcast_expand` no longer squeezes out batch sizes of size 1, as a batch size of 1 is still a
   batch size.
   [(#6147)](https://github.com/PennyLaneAI/pennylane/pull/6147)

pennylane/transforms/dynamic_one_shot.py

@@ -372,20 +372,21 @@ def gather_non_mcm(measurement, samples, is_valid, postselect_mode=None):
         tmp = Counter()
         for i, d in enumerate(samples):
             tmp.update(
-                dict((k if isinstance(k, str) else float(k), v * is_valid[i]) for k, v in d.items())
+                {k if isinstance(k, str) else float(k): v * is_valid[i] for k, v in d.items()}
             )
+
         if not measurement.all_outcomes:
             tmp = Counter({k: v for k, v in tmp.items() if v > 0})
         return dict(sorted(tmp.items()))
 
     if isinstance(measurement, SampleMP):
         if postselect_mode == "pad-invalid-samples" and samples.ndim == 2:
             is_valid = qml.math.reshape(is_valid, (-1, 1))
-        return (
-            qml.math.where(is_valid, samples, fill_in_value)
-            if postselect_mode == "pad-invalid-samples"
-            else samples[is_valid]
-        )
+        if postselect_mode == "pad-invalid-samples":
+            return qml.math.where(is_valid, samples, fill_in_value)
+        if qml.math.shape(samples) == ():  # single shot case
+            samples = qml.math.reshape(samples, (-1, 1))
+        return samples[is_valid]
 
     if (interface := qml.math.get_interface(is_valid)) == "tensorflow":
         # Tensorflow requires arrays that are used for arithmetic with each other to have the

tests/transforms/test_dynamic_one_shot.py

@@ -246,6 +246,7 @@ def generate_dummy_raw_results(measure_f, n_mcms, shots, postselect, interface):
         raw_results = (single_shot_res,) * shots
 
     else:
+
         # When postselecting, we start by creating results for two shots as alternating indices
         # will have valid results.
         # Alternating tuple. Only the values at odd indices are valid
@@ -276,7 +277,7 @@ class TestInterfaces:
     """Unit tests for ML interfaces with dynamic_one_shot"""
 
     @pytest.mark.parametrize("measure_f", (qml.expval, qml.probs, qml.sample, qml.var))
-    @pytest.mark.parametrize("shots", [20, [20, 21]])
+    @pytest.mark.parametrize("shots", [1, 20, [20, 21]])
     @pytest.mark.parametrize("n_mcms", [1, 3])
     def test_interface_tape_results(
         self, shots, n_mcms, measure_f, interface, use_interface_for_results
@@ -325,7 +326,7 @@ def test_interface_tape_results(
             (qml.var, 0.0, 0.0),
         ],
     )
-    @pytest.mark.parametrize("shots", [20, [20, 21]])
+    @pytest.mark.parametrize("shots", [1, 20, [20, 21]])
     @pytest.mark.parametrize("n_mcms", [1, 3])
     def test_interface_results_processing(
         self, shots, n_mcms, measure_f, expected1, expected2, interface, use_interface_for_results
@@ -394,7 +395,7 @@ def test_interface_results_processing(
             (qml.var, 0.0, 0.0),
         ],
     )
-    @pytest.mark.parametrize("shots", [20, [20, 22]])
+    @pytest.mark.parametrize("shots", [1, 20, [20, 22]])
     def test_interface_results_postselection_processing(
         self, shots, measure_f, expected1, expected2, interface, use_interface_for_results
     ):