Modify run_batch_async (#6387)

* Created using Colaboratory * call run_sweep_async in parallel * call circuits asynchronously * Revert "Created using Colaboratory" This reverts commit eb10318. * Revert "call run_sweep_async in parallel" This reverts commit f4e0d88. * revert colab * lint * use pmap * add test * lint * lint * nits
quantumlib · Dec 20, 2023 · 7c9b713 · 7c9b713
1 parent 99993b9
commit 7c9b713
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 4 deletions.
diff --git a/cirq-core/cirq/work/sampler.py b/cirq-core/cirq/work/sampler.py
@@ -294,10 +294,9 @@ async def run_batch_async(
         See docs for `cirq.Sampler.run_batch`.
         """
         params_list, repetitions = self._normalize_batch_args(programs, params_list, repetitions)
-        return [
-            await self.run_sweep_async(circuit, params=params, repetitions=repetitions)
-            for circuit, params, repetitions in zip(programs, params_list, repetitions)
-        ]
+        return await duet.pstarmap_async(
+            self.run_sweep_async, zip(programs, params_list, repetitions)
+        )
 
     def _normalize_batch_args(
         self,

diff --git a/cirq-core/cirq/work/sampler_test.py b/cirq-core/cirq/work/sampler_test.py
@@ -211,6 +211,31 @@ def test_sampler_run_batch():
         assert np.array_equal(result.measurements['m'], np.array([[0], [0]], dtype='uint8'))
 
 
+@duet.sync
+async def test_run_batch_async_calls_run_sweep_asynchronously():
+    """Test run_batch_async calls run_sweep_async without waiting."""
+    finished = []
+    a = cirq.LineQubit(0)
+    circuit1 = cirq.Circuit(cirq.X(a) ** sympy.Symbol('t'), cirq.measure(a, key='m'))
+    circuit2 = cirq.Circuit(cirq.Y(a) ** sympy.Symbol('t'), cirq.measure(a, key='m'))
+    params1 = cirq.Points('t', [0.3, 0.7])
+    params2 = cirq.Points('t', [0.4, 0.6])
+    params_list = [params1, params2]
+
+    class AsyncSampler(cirq.Sampler):
+        async def run_sweep_async(self, program, params, repetitions: int = 1):
+            if params == params1:
+                await duet.sleep(0.001)
+
+            finished.append(params)
+
+    await AsyncSampler().run_batch_async(
+        [circuit1, circuit2], params_list=params_list, repetitions=[1, 2]
+    )
+
+    assert finished == list(reversed(params_list))
+
+
 def test_sampler_run_batch_default_params_and_repetitions():
     sampler = cirq.ZerosSampler()
     a = cirq.LineQubit(0)