Use run_batch in cirq.Sampler (#678)

tensorflow · Apr 7, 2022 · 5ddde16 · 5ddde16
1 parent 4ef56df
commit 5ddde16
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/tensorflow_quantum/core/ops/cirq_ops.py b/tensorflow_quantum/core/ops/cirq_ops.py
@@ -529,10 +529,11 @@ def _no_grad(grad):
 
         else:
             # All other cirq.Samplers handled here.
-            #TODO(zaqqwerty): replace with run_batch once Cirq #3148 is resolved
             cirq_results = []
-            for p, r in zip(programs, resolvers):
-                cirq_results.append(sampler.run(p, r, num_samples))
+            for results in sampler.run_batch(programs,
+                                             params_list=resolvers,
+                                             repetitions=num_samples):
+                cirq_results.extend(results)
 
         results = []
         for r in cirq_results: