fix usage of memory-maps

Author: BatmanAoD

pyquil_for_azure_quantum/__init__.py

@@ -32,7 +32,11 @@
 from numpy import split
 from pyquil.api import QAM, MemoryMap, QAMExecutionResult, QuantumComputer, get_qc
 from pyquil.quil import Program
-from qcs_sdk import ExecutionData, RegisterData, ResultData  # pylint: disable=no-name-in-module
+from qcs_sdk import (  # pylint: disable=no-name-in-module
+    ExecutionData,
+    RegisterData,
+    ResultData,
+)
 from qcs_sdk.qvm import QVMResultData  # pylint: disable=no-name-in-module
 from wrapt import ObjectProxy
 
@@ -106,7 +110,7 @@ def compile(
     def run_batch(
         self,
         executable: AzureProgram,
-        memory_map: MemoryMap,
+        memory_maps: List[MemoryMap],
         **__kwargs: Any,
     ) -> List[QAMExecutionResult]:
         """Run a sequence of memory values through the program.
@@ -115,7 +119,7 @@ def run_batch(
             * [`AzureQuantumMachine.execute_with_memory_map_batch`][pyquil_for_azure_quantum.AzureQuantumMachine.execute_with_memory_map_batch]
         """
         qam = cast(AzureQuantumMachine, self.qam)
-        return qam.execute_with_memory_map_batch(executable, [memory_map])
+        return qam.execute_with_memory_map_batch(executable, memory_maps)
 
 
 def get_qpu(qpu_name: str) -> AzureQuantumComputer:
@@ -262,7 +266,7 @@ def get_result(self, execute_response: AzureJob) -> QAMExecutionResult:
             data=data,
         )
 
-    def execute_with_memory_map_batch(  # type: ignore
+    def execute_with_memory_map_batch(
         self,
         executable: AzureProgram,
         memory_maps: Iterable[MemoryMap],
@@ -273,14 +277,14 @@ def execute_with_memory_map_batch(  # type: ignore
 
         Args:
             executable: The AzureProgram to run.
-            memory_map: An iterable containing a ``MemoryMaps`` with desired mapping of parameter names to lists of parameter values.
-                Each value is a list as long as the number of slots in the register. So if the register was ``DECLARE theta REAL[2]``
-                then the key in the dictionary would be ``theta`` and the value would be a list of lists of length 2. The entire program
-                will be run (for shot count) as many times as there are values in the list. **All values (outer lists) must be of the same length**.
+            memory_maps: An iterable containing ``MemoryMaps`` with desired mappings of parameter names to parameter values.
+                Each value is a list as long as the number of elements in the register. So if the register was ``DECLARE theta REAL[2]``
+                then the key in the dictionary would be ``theta`` and the value would be a list of length 2. The entire program
+                will be run (for shot count) once for each ``MemoryMap``. (If no memory maps are provided, the program will be run once.)
             name: An optional name for the job which will show up in the Azure Quantum UI.
 
         Returns:
-            A list of ``QAMExecutionResult`` objects, one for each set of parameters.
+            A list of ``QAMExecutionResult`` objects, one for each ``MemoryMap``.
 
         ```pycon
 
@@ -309,48 +313,36 @@ def execute_with_memory_map_batch(  # type: ignore
 
         ```
         """
-        results = []
-        for memory_map in memory_maps:
-            num_params = None
-            for param_name, param_values in memory_map.items():
-                if num_params is None:
-                    num_params = len(param_values)
-                elif num_params != len(param_values):
-                    raise ValueError(
-                        "All parameter values must be of the same length. "
-                        f"{param_name} has length {len(param_values)} but {num_params} were expected."
+        executable = executable.copy()
+        input_params = InputParams(
+            count=executable.num_shots,
+            skip_quilc=executable.skip_quilc,
+            substitutions=memory_maps,
+        )
+        job = self._target.submit(
+            str(executable),
+            name=name,
+            input_params=input_params,
+        )
+        azure_job = AzureJob(job=job, executable=executable)
+        combined_result = self.get_result(azure_job)
+        num_executions = len(memory_maps)
+
+        if num_executions in (0, 1):
+            return combined_result
+
+        ro_matrix = combined_result.data.result_data.to_register_map().get_register_matrix("ro")
+        if ro_matrix is None:
+            return []
+
+        return [
+            QAMExecutionResult(
+                executable,
+                ExecutionData(
+                    ResultData.from_qvm(
+                        QVMResultData.from_memory_map(memory={"ro": RegisterData(split_result.tolist())})
                     )
-
-            executable = executable.copy()
-            input_params = InputParams(
-                count=executable.num_shots,
-                skip_quilc=executable.skip_quilc,
-                substitutions=memory_map,
-            )
-            job = self._target.submit(
-                str(executable),
-                name=name,
-                input_params=input_params,
+                ),
             )
-            azure_job = AzureJob(job=job, executable=executable)
-            combined_result = self.get_result(azure_job)
-            if num_params is None or num_params == 1:
-                results.append(combined_result)
-                continue
-
-            ro_matrix = combined_result.data.result_data.to_register_map().get_register_matrix("ro")
-            if ro_matrix is None:
-                continue
-
-            results += [
-                QAMExecutionResult(
-                    executable,
-                    ExecutionData(
-                        ResultData.from_qvm(
-                            QVMResultData.from_memory_map(memory={"ro": RegisterData(split_result.tolist())})
-                        )
-                    ),
-                )
-                for split_result in split(ro_matrix.to_ndarray(), num_params)
-            ]
-        return results
+            for split_result in split(ro_matrix.to_ndarray(), num_executions)
+        ]

test/test_e2e_no_qcs.py

@@ -72,8 +72,8 @@ def test_run_batch(qc: AzureQuantumComputer) -> None:
     """Test the ``run_batch`` interface which should be much faster than normal parametrization"""
     compiled = qc.compile(PARAMETRIZED)
 
-    memory_map: MemoryMap = {"theta": [0, np.pi, 2 * np.pi]}
-    results = qc.run_batch(compiled, memory_map)
+    executions = [{"theta": [value]} for value in [0, np.pi, 2 * np.pi]]
+    results = qc.run_batch(compiled, executions)
 
     results_0 = results[0].get_register_map().get("ro")
     assert results_0 is not None