Modify run_batch_async (quantumlib#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

Author: senecameeks

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,

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)