Return a single results object instead of always a list - IonQ

cirq-ionq/cirq_ionq/job.py

@@ -198,7 +198,12 @@ def results(
         polling_seconds: int = 1,
         sharpen: bool | None = None,
         extra_query_params: dict | None = None,
-    ) -> list[results.QPUResult] | list[results.SimulatorResult]:
+    ) -> (
+        results.QPUResult
+        | results.SimulatorResult
+        | list[results.QPUResult]
+        | list[results.SimulatorResult]
+    ):
         """Polls the IonQ api for results.
 
         Args:
@@ -209,16 +214,27 @@ def results(
             extra_query_params: Specify any parameters to include in the request.
 
         Returns:
-            Either a list of `cirq_ionq.QPUResult` or a list of `cirq_ionq.SimulatorResult`
-            depending on whether the job was running on an actual quantum processor or a
-            simulator.
+            Either a single `cirq_ionq.QPUResult` / `cirq_ionq.SimulatorResult`
+            (for a single-circuit job) or a `list` of such results (for a
+            batch job). The list order for batch jobs corresponds to the
+            order of the input circuits.
 
         Raises:
             IonQUnsuccessfulJob: If the job has failed, been canceled, or deleted.
             IonQException: If unable to get the results from the API.
             RuntimeError: If the job reported that it had failed on the server, or
                 the job had an unknown status.
             TimeoutError: If the job timed out at the server.
+
+        Notes:
+            * IonQ returns results in little endian; Cirq presents them in
+            big endian.
+            * If your code previously assumed a list, use:
+                r = job.results()
+                results_list = r if isinstance(r, list) else [r]
+            If your code previously assumed a single result, use:
+                r = job.results()
+                r0 = r[0] if isinstance(r, list) else r
         """
         time_waited_seconds = 0
         while time_waited_seconds < timeout_seconds:
@@ -245,11 +261,10 @@ def results(
             job_id=self.job_id(), sharpen=sharpen, extra_query_params=extra_query_params
         )
 
+        # is this a batch run (dict-of-dicts) or a single circuit?
         some_inner_value = next(iter(backend_results.values()))
-        if isinstance(some_inner_value, dict):
-            histograms = backend_results.values()
-        else:
-            histograms = [backend_results]
+        is_batch = isinstance(some_inner_value, dict)
+        histograms = list(backend_results.values()) if is_batch else [backend_results]
 
         # IonQ returns results in little endian, but
         # Cirq prefers to use big endian, so we convert.
@@ -270,7 +285,11 @@ def results(
                         measurement_dict=self.measurement_dict(circuit_index=circuit_index),
                     )
                 )
-            return big_endian_results_qpu
+            return (
+                big_endian_results_qpu
+                if len(big_endian_results_qpu) > 1
+                else big_endian_results_qpu[0]
+            )
         else:
             big_endian_results_sim: list[results.SimulatorResult] = []
             for circuit_index, histogram in enumerate(histograms):
@@ -286,7 +305,11 @@ def results(
                         repetitions=self.repetitions(),
                     )
                 )
-            return big_endian_results_sim
+            return (
+                big_endian_results_sim
+                if len(big_endian_results_sim) > 1
+                else big_endian_results_sim[0]
+            )
 
     def cancel(self):
         """Cancel the given job.

cirq-ionq/cirq_ionq/job_test.py

@@ -99,7 +99,7 @@ def test_job_results_qpu():
         assert "foo" in str(w[0].message)
         assert "bar" in str(w[1].message)
     expected = ionq.QPUResult({0: 600, 1: 400}, 2, {'a': [0, 1]})
-    assert results[0] == expected
+    assert results == expected
 
 
 def test_batch_job_results_qpu():
@@ -148,7 +148,7 @@ def test_job_results_rounding_qpu():
     job = ionq.Job(mock_client, job_dict)
     expected = ionq.QPUResult({0: 3, 1: 4997}, 2, {'a': [0, 1]})
     results = job.results()
-    assert results[0] == expected
+    assert results == expected
 
 
 def test_job_results_failed():
@@ -179,7 +179,7 @@ def test_job_results_qpu_endianness():
     }
     job = ionq.Job(mock_client, job_dict)
     results = job.results()
-    assert results[0] == ionq.QPUResult({0: 600, 2: 400}, 2, measurement_dict={})
+    assert results == ionq.QPUResult({0: 600, 2: 400}, 2, measurement_dict={})
 
 
 def test_batch_job_results_qpu_endianness():
@@ -200,7 +200,7 @@ def test_batch_job_results_qpu_endianness():
     }
     job = ionq.Job(mock_client, job_dict)
     results = job.results()
-    assert results[0] == ionq.QPUResult({0: 600, 2: 400}, 2, measurement_dict={'a': [0, 1]})
+    assert results == ionq.QPUResult({0: 600, 2: 400}, 2, measurement_dict={'a': [0, 1]})
 
 
 def test_job_results_qpu_target_endianness():
@@ -216,7 +216,7 @@ def test_job_results_qpu_target_endianness():
     }
     job = ionq.Job(mock_client, job_dict)
     results = job.results()
-    assert results[0] == ionq.QPUResult({0: 600, 2: 400}, 2, measurement_dict={})
+    assert results == ionq.QPUResult({0: 600, 2: 400}, 2, measurement_dict={})
 
 
 def test_batch_job_results_qpu_target_endianness():
@@ -238,7 +238,7 @@ def test_batch_job_results_qpu_target_endianness():
     }
     job = ionq.Job(mock_client, job_dict)
     results = job.results()
-    assert results[0] == ionq.QPUResult({0: 600, 2: 400}, 2, measurement_dict={'a': [0, 1]})
+    assert results == ionq.QPUResult({0: 600, 2: 400}, 2, measurement_dict={'a': [0, 1]})
 
 
 @mock.patch('time.sleep', return_value=None)
@@ -256,7 +256,7 @@ def test_job_results_poll(mock_sleep):
     mock_client.get_results.return_value = {'0': '0.6', '1': '0.4'}
     job = ionq.Job(mock_client, ready_job)
     results = job.results(polling_seconds=0)
-    assert results[0] == ionq.QPUResult({0: 600, 1: 400}, 1, measurement_dict={})
+    assert results == ionq.QPUResult({0: 600, 1: 400}, 1, measurement_dict={})
     mock_sleep.assert_called_once()
 
 
@@ -294,7 +294,7 @@ def test_job_results_simulator():
     }
     job = ionq.Job(mock_client, job_dict)
     results = job.results()
-    assert results[0] == ionq.SimulatorResult({0: 0.6, 1: 0.4}, 1, {}, 100)
+    assert results == ionq.SimulatorResult({0: 0.6, 1: 0.4}, 1, {}, 100)
 
 
 def test_batch_job_results_simulator():
@@ -336,7 +336,7 @@ def test_job_results_simulator_endianness():
     }
     job = ionq.Job(mock_client, job_dict)
     results = job.results()
-    assert results[0] == ionq.SimulatorResult({0: 0.6, 2: 0.4}, 2, {}, 100)
+    assert results == ionq.SimulatorResult({0: 0.6, 2: 0.4}, 2, {}, 100)
 
 
 def test_batch_job_results_simulator_endianness():
@@ -357,7 +357,7 @@ def test_batch_job_results_simulator_endianness():
     }
     job = ionq.Job(mock_client, job_dict)
     results = job.results()
-    assert results[0] == ionq.SimulatorResult({0: 0.6, 2: 0.4}, 2, {'a': [0, 1]}, 1000)
+    assert results == ionq.SimulatorResult({0: 0.6, 2: 0.4}, 2, {'a': [0, 1]}, 1000)
 
 
 def test_job_sharpen_results():
@@ -372,7 +372,7 @@ def test_job_sharpen_results():
     }
     job = ionq.Job(mock_client, job_dict)
     results = job.results(sharpen=False)
-    assert results[0] == ionq.SimulatorResult({0: 60, 1: 40}, 1, {}, 100)
+    assert results == ionq.SimulatorResult({0: 60, 1: 40}, 1, {}, 100)
 
 
 def test_job_cancel():

cirq-ionq/cirq_ionq/sampler.py

@@ -15,7 +15,6 @@
 
 from __future__ import annotations
 
-import itertools
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
@@ -107,7 +106,12 @@ def run_sweep(
             job_results = [job.results(timeout_seconds=self._timeout_seconds) for job in jobs]
         else:
             job_results = [job.results() for job in jobs]
-        flattened_job_results = list(itertools.chain.from_iterable(job_results))
+        flattened_job_results: list[results.QPUResult | results.SimulatorResult] = []
+        for res in job_results:
+            if isinstance(res, list):
+                flattened_job_results.extend(res)
+            else:
+                flattened_job_results.append(res)
         cirq_results = []
         for result, params in zip(flattened_job_results, resolvers):
             if isinstance(result, results.QPUResult):

cirq-ionq/cirq_ionq/sampler_test.py

@@ -154,3 +154,38 @@ def test_sampler_run_sweep():
         _, kwargs = call
         assert kwargs["repetitions"] == 4
         assert kwargs["target"] == 'qpu'
+
+
+def test_sampler_run_sweep_batched_job_results():
+    mock_service = mock.MagicMock()
+    job_dict = {
+        'id': '1',
+        'status': 'completed',
+        'stats': {'qubits': '1'},
+        'backend': 'qpu',
+        'metadata': {'shots': 4, 'measurement0': f'a{chr(31)}0'},
+    }
+
+    job = ionq.Job(client=mock_service, job_dict=job_dict)
+    mock_service.create_job.return_value = job
+
+    # Create dummy results
+    result1 = ionq.QPUResult(counts={0: 4}, num_qubits=1, measurement_dict={'a': [0]})
+    result2 = ionq.QPUResult(counts={1: 4}, num_qubits=1, measurement_dict={'a': [0]})
+
+    # Mock job.results to return a list
+    with mock.patch.object(ionq.Job, 'results', return_value=[result1, result2]):
+        sampler = ionq.Sampler(service=mock_service, target='qpu')
+        q0 = cirq.LineQubit(0)
+        circuit = cirq.Circuit(cirq.X(q0), cirq.measure(q0, key='a'))
+
+        # We pass 1 resolver, but job returns 2 results (list).
+        # This triggers `flattened_job_results.extend(res)`.
+        results = sampler.run_sweep(program=circuit, params=[cirq.ParamResolver({})], repetitions=4)
+
+        # Even though we got 2 results from the job, we only have 1 resolver,
+        # so zip() truncates to 1 result.
+        assert len(results) == 1
+        # result1 counts={0: 4}, so we expect four 0s.
+        # measurements is (repetitions, qubits), so [[0], [0], [0], [0]]
+        assert results[0].measurements['a'].tolist() == [[0], [0], [0], [0]]

cirq-ionq/cirq_ionq/service.py

@@ -157,13 +157,12 @@ def run(
             dry_run=dry_run,
             extra_query_params=extra_query_params,
         ).results(sharpen=sharpen)
-        if isinstance(job_results[0], results.QPUResult):
-            return job_results[0].to_cirq_result(params=cirq.ParamResolver(param_resolver))
-        if isinstance(job_results[0], results.SimulatorResult):
-            return job_results[0].to_cirq_result(
-                params=cirq.ParamResolver(param_resolver), seed=seed
-            )
-        raise NotImplementedError(f"Unrecognized job result type '{type(job_results[0])}'.")
+        result = job_results[0] if isinstance(job_results, list) else job_results
+        if isinstance(result, results.QPUResult):
+            return result.to_cirq_result(params=cirq.ParamResolver(param_resolver))
+        if isinstance(result, results.SimulatorResult):
+            return result.to_cirq_result(params=cirq.ParamResolver(param_resolver), seed=seed)
+        raise NotImplementedError(f"Unrecognized job result type '{type(result)}'.")
 
     def run_batch(
         self,
@@ -234,8 +233,9 @@ def run_batch(
             extra_query_params=extra_query_params,
         ).results(sharpen=sharpen)
 
+        job_results_list = job_results if isinstance(job_results, list) else [job_results]
         cirq_results = []
-        for job_result in job_results:
+        for job_result in job_results_list:
             if isinstance(job_result, results.QPUResult):
                 cirq_results.append(
                     job_result.to_cirq_result(params=cirq.ParamResolver(param_resolver))