From 5fbf1db01d4817b276b485960b4eb431cb399fc9 Mon Sep 17 00:00:00 2001 From: Cody Wang Date: Fri, 19 Jun 2026 16:59:17 -0700 Subject: [PATCH 1/3] feat: Automatically split/merge large program sets --- src/braket/pennylane_plugin/braket_device.py | 70 +++++---- test/unit_tests/test_braket_device.py | 154 +++++++++++-------- test/unit_tests/test_shadow_expval.py | 4 + 3 files changed, 133 insertions(+), 95 deletions(-) diff --git a/src/braket/pennylane_plugin/braket_device.py b/src/braket/pennylane_plugin/braket_device.py index d663775a..2e021eb1 100644 --- a/src/braket/pennylane_plugin/braket_device.py +++ b/src/braket/pennylane_plugin/braket_device.py @@ -84,7 +84,11 @@ ) from braket.program_sets import ProgramSet from braket.simulator import BraketSimulator -from braket.tasks import GateModelQuantumTaskResult, QuantumTask +from braket.tasks import ( + GateModelQuantumTaskResult, + ProgramSetQuantumTaskResult, + QuantumTask, +) from braket.tasks.local_quantum_task_batch import LocalQuantumTaskBatch from ._version import __version__ @@ -222,12 +226,6 @@ def batch_execute(self, circuits, **run_kwargs): if not self._parallel and not self._supports_program_sets: return super().batch_execute(circuits) - if self._supports_program_sets and ( - len(circuits) - > self._device.properties.action["braket.ir.openqasm.program_set"].maximumExecutables - ): - return super().batch_execute(circuits) - for circuit in circuits: self.check_validity(circuit.operations, circuit.observables) all_trainable = [] @@ -735,22 +733,41 @@ def use_grouping(self) -> bool: caps = self.capabilities() return not (caps.get("provides_jacobian")) + def _run_program_set( + self, program_set: ProgramSet, shots_per_executable: int + ) -> ProgramSetQuantumTaskResult: + program_sets, index_map = program_set.split( + self._device.properties.action[DeviceActionType.OPENQASM_PROGRAM_SET].maximumExecutables + ) + return ProgramSetQuantumTaskResult.merge( + [ + self._device.run( + sub_program_set, + s3_destination_folder=self._s3_folder, + shots=sub_program_set.total_executables * shots_per_executable, + poll_timeout_seconds=self._poll_timeout_seconds, + poll_interval_seconds=self._poll_interval_seconds, + **self._run_kwargs, + ).result() + for sub_program_set in program_sets + ], + program_set, + index_map, + ) + def _run_task_batch(self, braket_circuits, pl_circuits, batch_shots: int, inputs): if self._supports_program_sets: - program_set = ( - ProgramSet.zip(braket_circuits, input_sets=inputs) - if inputs - else ProgramSet(braket_circuits) - ) - task = self._device.run( - program_set, - s3_destination_folder=self._s3_folder, - shots=len(program_set) * batch_shots, - poll_timeout_seconds=self._poll_timeout_seconds, - poll_interval_seconds=self._poll_interval_seconds, - **self._run_kwargs, + return self._braket_program_set_to_pl_result( + self._run_program_set( + ( + ProgramSet.zip(braket_circuits, input_sets=inputs) + if inputs + else ProgramSet(braket_circuits) + ), + batch_shots, + ), + pl_circuits, ) - return self._braket_program_set_to_pl_result(task.result(), pl_circuits) task_batch = self._device.run_batch( braket_circuits, s3_destination_folder=self._s3_folder, @@ -794,16 +811,9 @@ def _run_snapshots(self, snapshot_circuits, n_qubits, mapped_wires): n_snapshots = len(snapshot_circuits) outcomes = np.zeros((n_snapshots, n_qubits)) if self._supports_program_sets: - program_set = ProgramSet(snapshot_circuits) - task = self._device.run( - program_set, - s3_destination_folder=self._s3_folder, - shots=len(program_set), - poll_timeout_seconds=self._poll_timeout_seconds, - poll_interval_seconds=self._poll_interval_seconds, - **self._run_kwargs, - ) - for t, result in enumerate(task.result()): + for t, result in enumerate( + self._run_program_set(ProgramSet(snapshot_circuits), shots_per_executable=1) + ): outcomes[t] = np.array(result[0].measurements[0])[mapped_wires] elif self._parallel: task_batch = self._device.run_batch( diff --git a/test/unit_tests/test_braket_device.py b/test/unit_tests/test_braket_device.py index eca7c65b..6f306000 100644 --- a/test/unit_tests/test_braket_device.py +++ b/test/unit_tests/test_braket_device.py @@ -225,6 +225,37 @@ ) +def _make_program_set_result(num_programs): + """Builds a ProgramSetQuantumTaskResult with ``num_programs`` single-executable programs, + each returning the measurements in PROGRAM_RESULT.""" + return ProgramSetQuantumTaskResult.from_object( + ProgramSetTaskResult( + **{ + "braketSchemaHeader": { + "name": "braket.task_result.program_set_task_result", + "version": "1", + }, + "programResults": [PROGRAM_RESULT] * num_programs, + "taskMetadata": { + "braketSchemaHeader": { + "name": "braket.task_result.program_set_task_metadata", + "version": "1", + }, + "id": "arn:aws:braket:us-west-2:667256736152:quantum-task/bfebc86f-e4ed-4d6f-8131-addd1a49d6dc", # noqa + "deviceId": "arn:aws:braket:::device/quantum-simulator/amazon/sv1", + "requestedShots": 20 * num_programs, + "successfulShots": 20 * num_programs, + "programMetadata": [{"executables": [{}]} for _ in range(num_programs)], + "createdAt": "2024-10-15T19:06:58.986Z", + "endedAt": "2024-10-15T19:07:00.382Z", + "status": "COMPLETED", + "totalFailedExecutables": 0, + }, + } + ) + ) + + DEVICE_ARN = "baz" @@ -1161,71 +1192,28 @@ def test_batch_execute_program_set_noncommuting(): @patch.object(AwsDevice, "run") def test_batch_execute_program_set_exceeds_max_executables(mock_run): - """Test batch_execute falls back to individual programs when exceeding maximumExecutables""" - custom_result = GateModelQuantumTaskResult.from_string( - json.dumps( - { - "braketSchemaHeader": { - "name": "braket.task_result.gate_model_task_result", - "version": "1", - }, - "measurements": [[0, 0], [1, 1], [0, 1], [1, 0]], - "resultTypes": [ - { - "type": { - "observable": ["x", "y"], - "targets": [0, 1], - "type": "expectation", - }, - "value": 0.5, - }, - { - "type": { - "observable": ["z", "z"], - "targets": [0, 1], - "type": "expectation", - }, - "value": 0.5, - }, - ], - "measuredQubits": [0, 1], - "taskMetadata": { - "braketSchemaHeader": { - "name": "braket.task_result.task_metadata", - "version": "1", - }, - "id": "task_arn", - "shots": 10000, - "deviceId": "default", - }, - "additionalMetadata": { - "action": { - "braketSchemaHeader": { - "name": "braket.ir.openqasm.program", - "version": "1", - }, - "source": "qubit[2] q; h q[0]; cnot q[0], q[1]; measure q;", - }, - }, - } - ) - ) + """Test batch_execute splits the program set and merges the results when the number of + executables exceeds the device's maximumExecutables.""" - # Mock the task to return our custom result - task = Mock() - task.result.return_value = custom_result - type(task).id = PropertyMock(return_value="task_arn") - task.state.return_value = "COMPLETED" - mock_run.return_value = task + # The program set is split into one task per chunk of maximumExecutables circuits, so each + # call to run returns a result sized to the chunk it was given. + def run_side_effect(program_set, **kwargs): + task = Mock() + task.result.return_value = _make_program_set_result(program_set.total_executables) + return task - dev = _aws_device(wires=4, foo="bar", parallel=True, supports_program_sets=True) + mock_run.side_effect = run_side_effect + + dev = _aws_device(wires=4, foo="bar", parallel=False, supports_program_sets=True) - # Verify the device properties are set correctly to ensure we hit the second condition - assert dev._parallel == True assert dev._supports_program_sets == True - assert dev._device.properties.action["braket.ir.openqasm.program_set"].maximumExecutables == 100 + max_executables = dev._device.properties.action[ + "braket.ir.openqasm.program_set" + ].maximumExecutables + assert max_executables == 100 - # Create 101 circuits (exceeds maximumExecutables of 100, defined in ACTION_PROPERTIES_PROGRAMSET) + # Create 101 circuits, exceeding maximumExecutables of 100 (defined in + # ACTION_PROPERTIES_PROGRAMSET), so the program set is split into two tasks. circuits = [] for _ in range(101): with QuantumTape() as circuit: @@ -1235,16 +1223,52 @@ def test_batch_execute_program_set_exceeds_max_executables(mock_run): circuits.append(circuit) assert len(circuits) == 101 - assert ( - len(circuits) - > dev._device.properties.action["braket.ir.openqasm.program_set"].maximumExecutables - ) + assert len(circuits) > max_executables result = dev.batch_execute(circuits) - assert mock_run.call_count == 101 + + # Two tasks: 100 executables in the first, 1 in the second. + assert mock_run.call_count == 2 + run_sizes = sorted(call.args[0].total_executables for call in mock_run.call_args_list) + assert run_sizes == [1, 100] + + # The merged result preserves the shape of the original (unsplit) batch. assert len(result) == 101 +@patch.object(AwsDevice, "run") +def test_run_snapshots_program_set_exceeds_max_executables(mock_run): + """Test _run_snapshots splits the program set and merges the results when the number of + snapshots exceeds the device's maximumExecutables.""" + + def run_side_effect(program_set, **kwargs): + task = Mock() + task.result.return_value = _make_program_set_result(program_set.total_executables) + return task + + mock_run.side_effect = run_side_effect + + dev = _aws_device(wires=2, foo="bar", parallel=False, supports_program_sets=True) + max_executables = dev._device.properties.action[ + "braket.ir.openqasm.program_set" + ].maximumExecutables + assert max_executables == 100 + + # 101 snapshots exceeds maximumExecutables of 100, so the program set is split into two tasks. + n_snapshots = 101 + snapshot_circuits = [Circuit().h(0).cnot(0, 1) for _ in range(n_snapshots)] + mapped_wires = np.array([0, 1]) + + outcomes = dev._run_snapshots(snapshot_circuits, n_qubits=2, mapped_wires=mapped_wires) + + assert mock_run.call_count == 2 + run_sizes = sorted(call.args[0].total_executables for call in mock_run.call_args_list) + assert run_sizes == [1, 100] + + # One outcome per snapshot, each shots=1 measurement projected onto the mapped wires. + assert outcomes.shape == (n_snapshots, 2) + + @patch.object(AwsDevice, "properties", new_callable=mock.PropertyMock) @patch.object(AwsDevice, "run_batch") def test_aws_device_batch_execute_parallel(mock_run_batch, mock_properties): diff --git a/test/unit_tests/test_shadow_expval.py b/test/unit_tests/test_shadow_expval.py index 84a395bd..f1579a48 100644 --- a/test/unit_tests/test_shadow_expval.py +++ b/test/unit_tests/test_shadow_expval.py @@ -374,6 +374,10 @@ def test_shadow_expval_aws_device( ): mock_action = Mock() mock_action.action = {"braket.ir.openqasm.program": None} + if supports_program_sets: + program_set_action = Mock() + program_set_action.maximumExecutables = 100 + mock_action.action[DeviceActionType.OPENQASM_PROGRAM_SET] = program_set_action mock_properties.return_value = mock_action dev = _aws_device( wires=2, From 4403de562d242f4957c4fa3cf1e412ca1f6b0525 Mon Sep 17 00:00:00 2001 From: Cody Wang Date: Wed, 8 Jul 2026 14:25:23 -0700 Subject: [PATCH 2/3] Run split program sets with `run_batch` --- src/braket/pennylane_plugin/braket_device.py | 48 +++--- test/unit_tests/test_braket_device.py | 166 +++++++++++++++---- test/unit_tests/test_shadow_expval.py | 2 +- 3 files changed, 159 insertions(+), 57 deletions(-) diff --git a/src/braket/pennylane_plugin/braket_device.py b/src/braket/pennylane_plugin/braket_device.py index 2e021eb1..a97e8c48 100644 --- a/src/braket/pennylane_plugin/braket_device.py +++ b/src/braket/pennylane_plugin/braket_device.py @@ -733,41 +733,47 @@ def use_grouping(self) -> bool: caps = self.capabilities() return not (caps.get("provides_jacobian")) - def _run_program_set( - self, program_set: ProgramSet, shots_per_executable: int - ) -> ProgramSetQuantumTaskResult: + def _run_program_set(self, program_set: ProgramSet) -> ProgramSetQuantumTaskResult: program_sets, index_map = program_set.split( self._device.properties.action[DeviceActionType.OPENQASM_PROGRAM_SET].maximumExecutables ) - return ProgramSetQuantumTaskResult.merge( + results = ( [ self._device.run( - sub_program_set, + program_sets[0], s3_destination_folder=self._s3_folder, - shots=sub_program_set.total_executables * shots_per_executable, + shots=program_sets[0].total_shots, poll_timeout_seconds=self._poll_timeout_seconds, poll_interval_seconds=self._poll_interval_seconds, **self._run_kwargs, ).result() - for sub_program_set in program_sets - ], - program_set, - index_map, + ] + if len(program_sets) == 1 + else self._device.run_batch( + program_sets, + s3_destination_folder=self._s3_folder, + shots=AwsDevice.DEFAULT_SHOTS_PROGRAM_SET, + max_parallel=self._max_parallel, + max_connections=self._max_connections, + poll_timeout_seconds=self._poll_timeout_seconds, + poll_interval_seconds=self._poll_interval_seconds, + **self._run_kwargs, + ).results(fail_unsuccessful=True, max_retries=self._max_retries) ) + return ProgramSetQuantumTaskResult.merge(results, program_set, index_map) def _run_task_batch(self, braket_circuits, pl_circuits, batch_shots: int, inputs): if self._supports_program_sets: - return self._braket_program_set_to_pl_result( - self._run_program_set( - ( - ProgramSet.zip(braket_circuits, input_sets=inputs) - if inputs - else ProgramSet(braket_circuits) - ), - batch_shots, - ), - pl_circuits, + result = self._run_program_set( + ProgramSet.zip( + braket_circuits, + input_sets=inputs, + shots_per_executable=batch_shots, + ) + if inputs + else ProgramSet(braket_circuits, shots_per_executable=batch_shots) ) + return self._braket_program_set_to_pl_result(result, pl_circuits) task_batch = self._device.run_batch( braket_circuits, s3_destination_folder=self._s3_folder, @@ -812,7 +818,7 @@ def _run_snapshots(self, snapshot_circuits, n_qubits, mapped_wires): outcomes = np.zeros((n_snapshots, n_qubits)) if self._supports_program_sets: for t, result in enumerate( - self._run_program_set(ProgramSet(snapshot_circuits), shots_per_executable=1) + self._run_program_set(ProgramSet(snapshot_circuits, shots_per_executable=1)) ): outcomes[t] = np.array(result[0].measurements[0])[mapped_wires] elif self._parallel: diff --git a/test/unit_tests/test_braket_device.py b/test/unit_tests/test_braket_device.py index 6f306000..38753510 100644 --- a/test/unit_tests/test_braket_device.py +++ b/test/unit_tests/test_braket_device.py @@ -1114,7 +1114,7 @@ def test_batch_execute_program_set(mock_run): braket_circuit = Circuit().h(0).cnot(0, 1).ry(0, -anp.pi / 2).rx(1, anp.pi / 2).i(2).i(3) mock_run.assert_called_with( - ProgramSet([braket_circuit, braket_circuit]), + ProgramSet([braket_circuit, braket_circuit], shots_per_executable=SHOTS), s3_destination_folder=("foo", "bar"), shots=SHOTS * 2, poll_timeout_seconds=AwsQuantumTask.DEFAULT_RESULTS_POLL_TIMEOUT, @@ -1162,6 +1162,7 @@ def test_batch_execute_program_set_parametrize_differentiable(mock_run): ProgramSet.zip( [braket_circuit1, braket_circuit2], input_sets=[{"p_0": -anp.pi / 2, "p_1": anp.pi / 2}, {"p_0": 0.123}], + shots_per_executable=SHOTS, ), s3_destination_folder=("foo", "bar"), shots=SHOTS * 2, @@ -1190,19 +1191,21 @@ def test_batch_execute_program_set_noncommuting(): dev.batch_execute(circuits) -@patch.object(AwsDevice, "run") -def test_batch_execute_program_set_exceeds_max_executables(mock_run): - """Test batch_execute splits the program set and merges the results when the number of - executables exceeds the device's maximumExecutables.""" +def _program_set_run_batch_mock(program_sets, **kwargs): + """Builds a task batch whose results are one ProgramSetQuantumTaskResult per split program + set, each sized to that program set's number of executables.""" + task_batch = Mock() + task_batch.results.return_value = [ + _make_program_set_result(program_set.total_executables) for program_set in program_sets + ] + return task_batch - # The program set is split into one task per chunk of maximumExecutables circuits, so each - # call to run returns a result sized to the chunk it was given. - def run_side_effect(program_set, **kwargs): - task = Mock() - task.result.return_value = _make_program_set_result(program_set.total_executables) - return task - mock_run.side_effect = run_side_effect +@patch.object(AwsDevice, "run_batch") +def test_batch_execute_program_set_exceeds_max_executables(mock_run_batch): + """Test batch_execute splits the program set and runs it as a batch when the number of + executables exceeds the device's maximumExecutables.""" + mock_run_batch.side_effect = _program_set_run_batch_mock dev = _aws_device(wires=4, foo="bar", parallel=False, supports_program_sets=True) @@ -1227,46 +1230,139 @@ def run_side_effect(program_set, **kwargs): result = dev.batch_execute(circuits) - # Two tasks: 100 executables in the first, 1 in the second. - assert mock_run.call_count == 2 - run_sizes = sorted(call.args[0].total_executables for call in mock_run.call_args_list) - assert run_sizes == [1, 100] + # A single batch of two program sets: 100 executables in one, 1 in the other, each + # carrying shots_per_executable so the service computes its own total shots. + mock_run_batch.assert_called_once() + program_sets = mock_run_batch.call_args.args[0] + assert sorted(program_set.total_executables for program_set in program_sets) == [1, 100] + assert all(program_set.shots_per_executable == SHOTS for program_set in program_sets) + assert mock_run_batch.call_args.kwargs["shots"] == AwsDevice.DEFAULT_SHOTS_PROGRAM_SET # The merged result preserves the shape of the original (unsplit) batch. assert len(result) == 101 -@patch.object(AwsDevice, "run") -def test_run_snapshots_program_set_exceeds_max_executables(mock_run): - """Test _run_snapshots splits the program set and merges the results when the number of - snapshots exceeds the device's maximumExecutables.""" +def _make_program_set_result_with_measurements(measurements): + """Builds a ProgramSetQuantumTaskResult with one single-executable, single-shot program per + entry in ``measurements``, where ``measurements[i]`` is the measured bitstring for program i.""" + program_results = [ + { + "braketSchemaHeader": { + "name": "braket.task_result.program_result", + "version": "1", + }, + "executableResults": [ + { + "braketSchemaHeader": { + "name": "braket.task_result.program_set_executable_result", + "version": "1", + }, + "measurements": [measurement], + "measuredQubits": list(range(len(measurement))), + "inputsIndex": 0, + } + ], + "source": { + "braketSchemaHeader": { + "name": "braket.ir.openqasm.program", + "version": "1", + }, + "source": "OPENQASM 3.0;", + }, + "additionalMetadata": { + "simulatorMetadata": { + "braketSchemaHeader": { + "name": "braket.task_result.simulator_metadata", + "version": "1", + }, + "executionDuration": 50, + } + }, + } + for measurement in measurements + ] + num_programs = len(measurements) + return ProgramSetQuantumTaskResult.from_object( + ProgramSetTaskResult( + **{ + "braketSchemaHeader": { + "name": "braket.task_result.program_set_task_result", + "version": "1", + }, + "programResults": program_results, + "taskMetadata": { + "braketSchemaHeader": { + "name": "braket.task_result.program_set_task_metadata", + "version": "1", + }, + "id": "arn:aws:braket:us-west-2:667256736152:quantum-task/bfebc86f-e4ed-4d6f-8131-addd1a49d6dc", # noqa + "deviceId": "arn:aws:braket:::device/quantum-simulator/amazon/sv1", + "requestedShots": num_programs, + "successfulShots": num_programs, + "programMetadata": [{"executables": [{}]} for _ in range(num_programs)], + "createdAt": "2024-10-15T19:06:58.986Z", + "endedAt": "2024-10-15T19:07:00.382Z", + "status": "COMPLETED", + "totalFailedExecutables": 0, + }, + } + ) + ) - def run_side_effect(program_set, **kwargs): - task = Mock() - task.result.return_value = _make_program_set_result(program_set.total_executables) - return task - mock_run.side_effect = run_side_effect +@patch.object(AwsDevice, "run_batch") +def test_run_snapshots_program_set_exceeds_max_executables(mock_run_batch): + """Test _run_snapshots splits the program set and runs it as a batch when the number of + snapshots exceeds the device's maximumExecutables, and that outcomes are merged back in the + original snapshot order.""" + n_snapshots = 101 + n_qubits = 7 # enough bits to uniquely encode each snapshot index in 0..100 + + def measurement_for(index): + # Big-endian bit encoding of the snapshot index, so each snapshot has a unique measurement. + return [(index >> (n_qubits - 1 - bit)) & 1 for bit in range(n_qubits)] + + def run_batch_side_effect(program_sets, **kwargs): + # The split preserves order, so the executables across the returned program sets are a + # contiguous partition of the original snapshots. Encode each executable's original index + # into its measurement so the test can detect any reordering during merge. + task_batch = Mock() + results = [] + offset = 0 + for program_set in program_sets: + size = program_set.total_executables + results.append( + _make_program_set_result_with_measurements( + [measurement_for(offset + i) for i in range(size)] + ) + ) + offset += size + task_batch.results.return_value = results + return task_batch + + mock_run_batch.side_effect = run_batch_side_effect - dev = _aws_device(wires=2, foo="bar", parallel=False, supports_program_sets=True) + dev = _aws_device(wires=n_qubits, foo="bar", parallel=False, supports_program_sets=True) max_executables = dev._device.properties.action[ "braket.ir.openqasm.program_set" ].maximumExecutables assert max_executables == 100 # 101 snapshots exceeds maximumExecutables of 100, so the program set is split into two tasks. - n_snapshots = 101 snapshot_circuits = [Circuit().h(0).cnot(0, 1) for _ in range(n_snapshots)] - mapped_wires = np.array([0, 1]) - - outcomes = dev._run_snapshots(snapshot_circuits, n_qubits=2, mapped_wires=mapped_wires) + mapped_wires = np.arange(n_qubits) - assert mock_run.call_count == 2 - run_sizes = sorted(call.args[0].total_executables for call in mock_run.call_args_list) - assert run_sizes == [1, 100] + outcomes = dev._run_snapshots(snapshot_circuits, n_qubits=n_qubits, mapped_wires=mapped_wires) - # One outcome per snapshot, each shots=1 measurement projected onto the mapped wires. - assert outcomes.shape == (n_snapshots, 2) + mock_run_batch.assert_called_once() + program_sets = mock_run_batch.call_args.args[0] + assert sorted(program_set.total_executables for program_set in program_sets) == [1, 100] + + # One outcome per snapshot; each outcome must decode back to its original snapshot index, + # confirming the split results were merged back in the original order. + assert outcomes.shape == (n_snapshots, n_qubits) + expected = np.array([measurement_for(t) for t in range(n_snapshots)]) + assert np.array_equal(outcomes, expected) @patch.object(AwsDevice, "properties", new_callable=mock.PropertyMock) diff --git a/test/unit_tests/test_shadow_expval.py b/test/unit_tests/test_shadow_expval.py index f1579a48..825a62e2 100644 --- a/test/unit_tests/test_shadow_expval.py +++ b/test/unit_tests/test_shadow_expval.py @@ -354,7 +354,7 @@ def test_only_one_operator_in_shadow_expval(): [ (False, circs, TASK, SHOTS, None, None, False), (True, circs, TASK_BATCH, 1, 10, 10, False), - (False, ProgramSet(circs), TASK_PROGRAM_SET, 1, None, None, True), + (False, ProgramSet(circs, shots_per_executable=1), TASK_PROGRAM_SET, 1, None, None, True), ], ) def test_shadow_expval_aws_device( From b0690985acae1e9bc4f46036e18f68ce180d1e19 Mon Sep 17 00:00:00 2001 From: Cody Wang Date: Wed, 8 Jul 2026 14:53:25 -0700 Subject: [PATCH 3/3] Make max executables a device property --- src/braket/pennylane_plugin/braket_device.py | 20 ++++++++++---------- test/unit_tests/test_braket_device.py | 15 ++++----------- test/unit_tests/test_shadow_expval.py | 8 +++----- 3 files changed, 17 insertions(+), 26 deletions(-) diff --git a/src/braket/pennylane_plugin/braket_device.py b/src/braket/pennylane_plugin/braket_device.py index a97e8c48..f3bf7470 100644 --- a/src/braket/pennylane_plugin/braket_device.py +++ b/src/braket/pennylane_plugin/braket_device.py @@ -185,9 +185,11 @@ def __init__( self._supported_obs = supported_observables(self._device, self.shots) self._check_supported_result_types() self._verbatim = verbatim - self._supports_program_sets = ( - DeviceActionType.OPENQASM_PROGRAM_SET in self._device.properties.action + self._max_program_set_executables = ( + self._device.properties.action[DeviceActionType.OPENQASM_PROGRAM_SET].maximumExecutables + if DeviceActionType.OPENQASM_PROGRAM_SET in self._device.properties.action and self._shots is not None + else None ) if noise_model: @@ -223,7 +225,7 @@ def parallel(self) -> bool: return self._parallel def batch_execute(self, circuits, **run_kwargs): - if not self._parallel and not self._supports_program_sets: + if not self._parallel and self._max_program_set_executables is None: return super().batch_execute(circuits) for circuit in circuits: @@ -242,7 +244,7 @@ def batch_execute(self, circuits, **run_kwargs): self._pl_to_braket_circuit( circuit, trainable_indices=frozenset(trainable.keys()), - add_observables=not self._supports_program_sets, + add_observables=self._max_program_set_executables is None, **run_kwargs, ) ) @@ -734,9 +736,7 @@ def use_grouping(self) -> bool: return not (caps.get("provides_jacobian")) def _run_program_set(self, program_set: ProgramSet) -> ProgramSetQuantumTaskResult: - program_sets, index_map = program_set.split( - self._device.properties.action[DeviceActionType.OPENQASM_PROGRAM_SET].maximumExecutables - ) + program_sets, index_map = program_set.split(self._max_program_set_executables) results = ( [ self._device.run( @@ -763,7 +763,7 @@ def _run_program_set(self, program_set: ProgramSet) -> ProgramSetQuantumTaskResu return ProgramSetQuantumTaskResult.merge(results, program_set, index_map) def _run_task_batch(self, braket_circuits, pl_circuits, batch_shots: int, inputs): - if self._supports_program_sets: + if self._max_program_set_executables is not None: result = self._run_program_set( ProgramSet.zip( braket_circuits, @@ -816,7 +816,7 @@ def _run_task(self, circuit, inputs=None): def _run_snapshots(self, snapshot_circuits, n_qubits, mapped_wires): n_snapshots = len(snapshot_circuits) outcomes = np.zeros((n_snapshots, n_qubits)) - if self._supports_program_sets: + if self._max_program_set_executables is not None: for t, result in enumerate( self._run_program_set(ProgramSet(snapshot_circuits, shots_per_executable=1)) ): @@ -1160,7 +1160,7 @@ def __init__( super().__init__(wires, device, shots=shots, **run_kwargs) # TODO: Enable program sets once local simulator supports multiprocessing # for program set execution - self._supports_program_sets = False + self._max_program_set_executables = None def _run_task_batch(self, braket_circuits, pl_circuits, batch_shots: int, inputs): task_batch = self._device.run_batch( diff --git a/test/unit_tests/test_braket_device.py b/test/unit_tests/test_braket_device.py index 38753510..9195644b 100644 --- a/test/unit_tests/test_braket_device.py +++ b/test/unit_tests/test_braket_device.py @@ -1209,11 +1209,7 @@ def test_batch_execute_program_set_exceeds_max_executables(mock_run_batch): dev = _aws_device(wires=4, foo="bar", parallel=False, supports_program_sets=True) - assert dev._supports_program_sets == True - max_executables = dev._device.properties.action[ - "braket.ir.openqasm.program_set" - ].maximumExecutables - assert max_executables == 100 + assert dev._max_program_set_executables == 100 # Create 101 circuits, exceeding maximumExecutables of 100 (defined in # ACTION_PROPERTIES_PROGRAMSET), so the program set is split into two tasks. @@ -1226,7 +1222,7 @@ def test_batch_execute_program_set_exceeds_max_executables(mock_run_batch): circuits.append(circuit) assert len(circuits) == 101 - assert len(circuits) > max_executables + assert len(circuits) > dev._max_program_set_executables result = dev.batch_execute(circuits) @@ -1343,10 +1339,7 @@ def run_batch_side_effect(program_sets, **kwargs): mock_run_batch.side_effect = run_batch_side_effect dev = _aws_device(wires=n_qubits, foo="bar", parallel=False, supports_program_sets=True) - max_executables = dev._device.properties.action[ - "braket.ir.openqasm.program_set" - ].maximumExecutables - assert max_executables == 100 + assert dev._max_program_set_executables == 100 # 101 snapshots exceeds maximumExecutables of 100, so the program set is split into two tasks. snapshot_circuits = [Circuit().h(0).cnot(0, 1) for _ in range(n_snapshots)] @@ -1484,7 +1477,7 @@ def test_local_sim_batch_execute_parallel(mock_run_batch): RESULT.get_value_by_result_type(result_types.Sample(observable=observables.Z(3))), ) - if dev._supports_program_sets: + if dev._max_program_set_executables is not None: expected_circuits = [CIRCUIT_WITH_BASIS_ROTATION, CIRCUIT_WITH_BASIS_ROTATION] else: expected_circuits = [CIRCUIT_DIAGONALIZED, CIRCUIT_DIAGONALIZED] diff --git a/test/unit_tests/test_shadow_expval.py b/test/unit_tests/test_shadow_expval.py index 825a62e2..990f17ad 100644 --- a/test/unit_tests/test_shadow_expval.py +++ b/test/unit_tests/test_shadow_expval.py @@ -374,10 +374,6 @@ def test_shadow_expval_aws_device( ): mock_action = Mock() mock_action.action = {"braket.ir.openqasm.program": None} - if supports_program_sets: - program_set_action = Mock() - program_set_action.maximumExecutables = 100 - mock_action.action[DeviceActionType.OPENQASM_PROGRAM_SET] = program_set_action mock_properties.return_value = mock_action dev = _aws_device( wires=2, @@ -483,7 +479,9 @@ def _aws_device( ): properties_mock.action = {DeviceActionType.OPENQASM: action_properties} if supports_program_sets: - properties_mock.action[DeviceActionType.OPENQASM_PROGRAM_SET] = action_properties + program_set_action = Mock() + program_set_action.maximumExecutables = 100 + properties_mock.action[DeviceActionType.OPENQASM_PROGRAM_SET] = program_set_action type_mock.return_value = device_type dev = BraketAwsQubitDevice( wires=wires,