cirq-ops parallelization for non-QuantumEngine backends

Hi,

I work at IonQ on our quantum cloud, we recently started looking into our tensorflow quantum support and usability, and found that predict() and other high level tensorflow uses were blocking, despite efforts to use generator Datasets / setting workers, and doing our best to enable tfq parallelization.

There appears to be some custom code that enables parallelization for QuantumEngine here:

quantum/tensorflow_quantum/core/ops/cirq_ops.py

Lines 494 to 530 in 768ac15

    
           if isinstance(sampler, cirq_google.QuantumEngineSampler): 
        
               # group samples from identical circuits to reduce communication 
        
               # overhead. Have to keep track of the order in which things came 
        
               # in to make sure the output is ordered correctly 
        
               to_be_grouped = [ 
        
                   (ser_prog.numpy(), resolver, index) 
        
                   for index, ( 
        
                       ser_prog, 
        
                       resolver) in enumerate(zip(serialized_programs, resolvers)) 
        
               ] 
        
               grouped = _group_tuples(to_be_grouped) 
        
               # start all the necessary jobs 
        
               results_mapping = {} 
        
               for key, value in grouped.items(): 
        
                   program = programs[value[0][1]] 
        
                   resolvers = [x[0] for x in value] 
        
                   orders = [x[1] for x in value] 
        
                   # sampler.run_sweep blocks until results are in, so go around it 
        
                   result = sampler._engine.run_sweep( 
        
                       program=program, 
        
                       params=resolvers, 
        
                       repetitions=num_samples, 
        
                       processor_ids=sampler._processor_ids, 
        
                       gate_set=sampler._gate_set) 
        
                   results_mapping[result] = orders 
        
               # get all results 
        
               cirq_results = [None] * len(programs) 
        
               for key, value in results_mapping.items(): 
        
                   this_results = key.results() 
        
                   for result, index in zip(this_results, value): 
        
                       cirq_results[index] = result 
        
           else:

Whereas our code, falling through, hits the blocking path sampler.run:

quantum/tensorflow_quantum/core/ops/cirq_ops.py

Line 535 in 768ac15

cirq_results.append(sampler.run(p, r, num_samples))

Requesting guidance on the best path forward. We would like to ensure cirq_ops is not blocking for our Cirq backend. Should we look at doing something similar to the QuantumEngine override? Is there other work that we could help move along to make things non-blocking across backends? Thank you in advance!

Ah, the linked issue ties to quantumlib/Cirq#3224 which is fixed. Will look at whether changing to run_batch/sweep/etc would help us here as well

I think implementing run_batch works for us, closing this

	if isinstance(sampler, cirq_google.QuantumEngineSampler):
	# group samples from identical circuits to reduce communication
	# overhead. Have to keep track of the order in which things came
	# in to make sure the output is ordered correctly
	to_be_grouped = [
	(ser_prog.numpy(), resolver, index)
	for index, (
	ser_prog,
	resolver) in enumerate(zip(serialized_programs, resolvers))
	]

	grouped = _group_tuples(to_be_grouped)

	# start all the necessary jobs
	results_mapping = {}
	for key, value in grouped.items():
	program = programs[value[0][1]]
	resolvers = [x[0] for x in value]
	orders = [x[1] for x in value]

	# sampler.run_sweep blocks until results are in, so go around it
	result = sampler._engine.run_sweep(
	program=program,
	params=resolvers,
	repetitions=num_samples,
	processor_ids=sampler._processor_ids,
	gate_set=sampler._gate_set)
	results_mapping[result] = orders

	# get all results
	cirq_results = [None] * len(programs)
	for key, value in results_mapping.items():
	this_results = key.results()
	for result, index in zip(this_results, value):
	cirq_results[index] = result

	else: