diff --git a/setup.py b/setup.py index b586489f..801bfeb3 100644 --- a/setup.py +++ b/setup.py @@ -34,7 +34,7 @@ packages=find_namespace_packages(where="src", exclude=("test",)), package_dir={"": "src"}, install_requires=[ - "amazon-braket-sdk>=1.87.0", + "amazon-braket-sdk>=1.97.0", "autoray>=0.6.11", "pennylane>=0.34.0", ], diff --git a/src/braket/pennylane_plugin/braket_device.py b/src/braket/pennylane_plugin/braket_device.py index 24123a87..bdb38efe 100644 --- a/src/braket/pennylane_plugin/braket_device.py +++ b/src/braket/pennylane_plugin/braket_device.py @@ -83,13 +83,14 @@ translate_result, translate_result_type, ) +from braket.program_sets import ProgramSet from braket.simulator import BraketSimulator from braket.tasks import GateModelQuantumTaskResult, QuantumTask from braket.tasks.local_quantum_task_batch import LocalQuantumTaskBatch from ._version import __version__ -RETURN_TYPES = (ExpectationMP, VarianceMP, SampleMP, ProbabilityMP, StateMP, CountsMP) +RETURN_TYPES = (ExpectationMP, VarianceMP, SampleMP, ProbabilityMP, StateMP, CountsMP) MIN_SIMULATOR_BILLED_MS = 3000 OBS_LIST = (qml.PauliX, qml.PauliY, qml.PauliZ) @@ -168,6 +169,10 @@ 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 + and self._shots is not None + ) if noise_model: self._validate_noise_model_support() @@ -202,7 +207,7 @@ def parallel(self) -> bool: return self._parallel def batch_execute(self, circuits, **run_kwargs): - if not self._parallel: + if not self._parallel and not self._supports_program_sets: return super().batch_execute(circuits) for circuit in circuits: @@ -220,6 +225,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, **run_kwargs, ) ) @@ -232,18 +238,15 @@ def batch_execute(self, circuits, **run_kwargs): else [] ) - braket_results_batch = self._run_task_batch(braket_circuits, batch_shots, batch_inputs) - - return [ - self._braket_to_pl_result(braket_result, circuit) - for braket_result, circuit in zip(braket_results_batch, circuits) - ] + return self._run_task_batch(braket_circuits, circuits, batch_shots, batch_inputs) def _pl_to_braket_circuit( self, circuit: QuantumTape, compute_gradient: bool = False, trainable_indices: frozenset[int] = None, + *, + add_observables: bool = True, **run_kwargs, ): """Converts a PennyLane circuit to a Braket circuit""" @@ -259,17 +262,28 @@ def _pl_to_braket_circuit( if compute_gradient: braket_circuit = self._apply_gradient_result_type(circuit, braket_circuit) elif not isinstance(circuit.measurements[0], MeasurementTransform): - for measurement in circuit.measurements: - translated = translate_result_type( - measurement.map_wires(self.wire_map), - None, - self._braket_result_types, + if add_observables: + for measurement in circuit.measurements: + translated = translate_result_type( + measurement.map_wires(self.wire_map), + None, + self._braket_result_types, + ) + if isinstance(translated, tuple): + for result_type in translated: + braket_circuit.add_result_type(result_type) + else: + braket_circuit.add_result_type(translated) + else: + groups = qml.pauli.group_observables( + [measurement.obs for measurement in circuit.measurements], grouping_type="qwc" ) - if isinstance(translated, tuple): - for result_type in translated: - braket_circuit.add_result_type(result_type) - else: - braket_circuit.add_result_type(translated) + if len(groups) > 1: + raise ValueError( + f"Observables need to mutually commute, but found {len(groups)}: {groups}" + ) + diagonalizing_ops = qml.pauli.diagonalize_qwc_pauli_words(groups[0])[0] + braket_circuit += self.apply(diagonalizing_ops, apply_identities=False) return braket_circuit @@ -316,7 +330,7 @@ def _update_tracker_for_batch( self.tracker.update(batches=1, executions=total_executions, shots=total_shots) self.tracker.record() - def statistics( + def _statistics( self, braket_result: GateModelQuantumTaskResult, measurements: Sequence[MeasurementProcess], @@ -338,14 +352,18 @@ def statistics( for mp in measurements: if not isinstance(mp, RETURN_TYPES): raise QuantumFunctionError("Unsupported return type: {}".format(type(mp))) - results.append(self._get_statistic(braket_result, mp)) + results.append( + translate_result( + braket_result, mp.map_wires(self.wire_map), None, self._braket_result_types + ) + ) return results def _braket_to_pl_result(self, braket_result, circuit): """Calculates the PennyLane results from a Braket task result. A PennyLane circuit also determines the output observables.""" # Compute the required statistics - results = self.statistics(braket_result, circuit.measurements) + results = self._statistics(braket_result, circuit.measurements) ag_results = [ result for result in braket_result.result_types @@ -378,6 +396,25 @@ def _braket_to_pl_result(self, braket_result, circuit): return onp.array(results).squeeze() return tuple(onp.array(result).squeeze() for result in results) + def _braket_program_set_to_pl_result(self, program_set_result, circuits): + results = [] + for program_result, circuit in zip(program_set_result, circuits): + # Only one executable per program + measurements = program_result[0].measurements + + # Program sets require shots > 0, + # so the circuit's measurements are guaranteed to be SampleMeasurements + executable_results = [ + measurement.process_samples(measurements, wire_order=measurement.wires) + for measurement in circuit.measurements + ] + results.append( + onp.array(executable_results).squeeze() + if len(circuit.measurements) == 1 + else tuple(onp.array(result).squeeze() for result in executable_results) + ) + return results + @staticmethod def _tracking_data(task): if task.state() == "COMPLETED": @@ -410,8 +447,6 @@ def classical_shadow(self, obs, circuit): rng = np.random.default_rng(seed) recipes = rng.integers(0, 3, size=(n_snapshots, n_qubits)) - outcomes = np.zeros((n_snapshots, n_qubits)) - snapshot_rotations = [ [ rot @@ -484,6 +519,7 @@ def apply( use_unique_params: bool = False, *, trainable_indices: Optional[frozenset[int]] = None, + apply_identities: bool = True, **run_kwargs, ) -> Circuit: """Instantiate Braket Circuit object.""" @@ -518,8 +554,9 @@ def apply( unused = set(range(self.num_wires)) - {int(qubit) for qubit in circuit.qubits} # To ensure the results have the right number of qubits - for qubit in sorted(unused): - circuit.i(qubit) + if apply_identities: + for qubit in sorted(unused): + circuit.i(qubit) if self._noise_model: circuit = self._noise_model.apply(circuit) @@ -552,14 +589,12 @@ def _validate_noise_model_support(self): def _run_task(self, circuit, inputs=None): raise NotImplementedError("Need to implement task runner") + def _run_task_batch(self, braket_circuits, pl_circuits, circuit_shots, mapped_wires): + raise NotImplementedError("Need to implement batch runner") + def _run_snapshots(self, snapshot_circuits, n_qubits, mapped_wires): raise NotImplementedError("Need to implement snapshots runner") - def _get_statistic(self, braket_result, mp): - return translate_result( - braket_result, mp.map_wires(self.wire_map), None, self._braket_result_types - ) - @staticmethod def _get_trainable_parameters(tape: QuantumTape) -> dict[int, numbers.Number]: trainable_indices = sorted(tape.trainable_params) @@ -663,9 +698,24 @@ def use_grouping(self) -> bool: caps = self.capabilities() return not ("provides_jacobian" in caps and caps["provides_jacobian"]) - def _run_task_batch(self, batch_circuits, batch_shots: int, inputs): + 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(task.result(), pl_circuits) task_batch = self._device.run_batch( - batch_circuits, + braket_circuits, s3_destination_folder=self._s3_folder, shots=batch_shots, max_parallel=self._max_parallel, @@ -687,7 +737,10 @@ def _run_task_batch(self, batch_circuits, batch_shots: int, inputs): if self.tracker.active: self._update_tracker_for_batch(task_batch, batch_shots) - return braket_results_batch + return [ + self._braket_to_pl_result(braket_result, circuit) + for braket_result, circuit in zip(braket_results_batch, pl_circuits) + ] def _run_task(self, circuit, inputs=None): return self._device.run( @@ -703,7 +756,19 @@ 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._parallel: + 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()): + outcomes[t] = np.array(result[0].measurements[0])[mapped_wires] + elif self._parallel: task_batch = self._device.run_batch( snapshot_circuits, s3_destination_folder=self._s3_folder, @@ -1041,9 +1106,9 @@ def __init__( device = LocalSimulator(backend) super().__init__(wires, device, shots=shots, **run_kwargs) - def _run_task_batch(self, batch_circuits, batch_shots: int, inputs): + def _run_task_batch(self, braket_circuits, pl_circuits, batch_shots: int, inputs): task_batch = self._device.run_batch( - batch_circuits, + braket_circuits, shots=batch_shots, max_parallel=self._max_parallel, inputs=inputs, @@ -1057,7 +1122,10 @@ def _run_task_batch(self, batch_circuits, batch_shots: int, inputs): if self.tracker.active: self._update_tracker_for_batch(task_batch, batch_shots) - return braket_results_batch + return [ + self._braket_to_pl_result(braket_result, circuit) + for braket_result, circuit in zip(braket_results_batch, pl_circuits) + ] def _run_task(self, circuit, inputs=None): return self._device.run( diff --git a/src/braket/pennylane_plugin/translation.py b/src/braket/pennylane_plugin/translation.py index 3eb8a9ed..1ba946a6 100644 --- a/src/braket/pennylane_plugin/translation.py +++ b/src/braket/pennylane_plugin/translation.py @@ -594,7 +594,9 @@ def translate_result_type( # noqa: C901 if isinstance(observable, qml.ops.LinearCombination): if isinstance(measurement, qml.measurements.ExpectationMP): return tuple(Expectation(_translate_observable(op)) for op in observable.terms()[1]) - raise NotImplementedError(f"Return type {type(measurement)} unsupported for LinearCombination") + raise NotImplementedError( + f"Return type {type(measurement)} unsupported for LinearCombination" + ) braket_observable = _translate_observable(observable) if isinstance(measurement, qml.measurements.ExpectationMP): @@ -722,7 +724,11 @@ def translate_result( ] targets = targets or measurement.wires.tolist() - if isinstance(measurement, qml.measurements.CountsMP) and not measurement.all_outcomes and observable is None: + if ( + isinstance(measurement, qml.measurements.CountsMP) + and not measurement.all_outcomes + and observable is None + ): if targets: new_dict = {} for key, value in braket_result.measurement_counts.items(): diff --git a/test/unit_tests/test_braket_device.py b/test/unit_tests/test_braket_device.py index 789fa69f..954b81aa 100644 --- a/test/unit_tests/test_braket_device.py +++ b/test/unit_tests/test_braket_device.py @@ -13,7 +13,6 @@ import json from collections import Counter -from enum import Enum from typing import Any, Optional from unittest import mock from unittest.mock import Mock, PropertyMock, patch @@ -41,8 +40,10 @@ ) from braket.device_schema.simulators import GateModelSimulatorDeviceCapabilities from braket.devices import LocalSimulator +from braket.program_sets import ProgramSet from braket.simulator import BraketSimulator -from braket.tasks import GateModelQuantumTaskResult +from braket.task_result import ProgramSetTaskResult +from braket.tasks import GateModelQuantumTaskResult, ProgramSetQuantumTaskResult from device_property_jsons import ( ACTION_PROPERTIES, ACTION_PROPERTIES_DM_DEVICE, @@ -97,6 +98,104 @@ .sample(observable=observables.Z(3)) ) +CIRCUIT_BELL = Circuit().h(0).cnot(0, 1) +PROGRAM_RESULT = { + "braketSchemaHeader": { + "name": "braket.task_result.program_result", + "version": "1", + }, + "executableResults": [ + { + "braketSchemaHeader": { + "name": "braket.task_result.program_set_executable_result", + "version": "1", + }, + "measurements": [ + [0, 0], + [0, 1], + [1, 1], + [0, 0], + [1, 1], + [0, 0], + [1, 1], + [1, 0], + [1, 1], + [0, 0], + [1, 1], + [0, 0], + [0, 1], + [1, 0], + [1, 1], + [1, 1], + [1, 1], + [0, 0], + [1, 1], + [0, 0], + ], + "measuredQubits": [0, 1], + "inputsIndex": 0, + } + ], + "source": { + "braketSchemaHeader": { + "name": "braket.ir.openqasm.program", + "version": "1", + }, + "source": "OPENQASM 3.0;\nbit[2] b;\nqubit[2] q;\nh q[0];\ncnot q[0], q[1];\nb[0] = measure q[0];\nb[1] = measure q[1];", # noqa + "inputs": {"theta": [0.12, 2.1]}, + }, + "additionalMetadata": { + "simulatorMetadata": { + "braketSchemaHeader": { + "name": "braket.task_result.simulator_metadata", + "version": "1", + }, + "executionDuration": 50, + } + }, +} +PROGRAM_SET_RESULT = ProgramSetQuantumTaskResult.from_object( + ProgramSetTaskResult( + **{ + "braketSchemaHeader": { + "name": "braket.task_result.program_set_task_result", + "version": "1", + }, + "programResults": [PROGRAM_RESULT] * 2, + "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": 120, + "successfulShots": 100, + "programMetadata": [{"executables": [{}]}], + "deviceParameters": { + "braketSchemaHeader": { + "name": "braket.device_schema.simulators.gate_model_simulator_device_parameters", + "version": "1", + }, + "paradigmParameters": { + "braketSchemaHeader": { + "name": "braket.device_schema.gate_model_parameters", + "version": "1", + }, + "qubitCount": 5, + "disableQubitRewiring": False, + }, + }, + "createdAt": "2024-10-15T19:06:58.986Z", + "endedAt": "2024-10-15T19:07:00.382Z", + "status": "COMPLETED", + "totalFailedExecutables": 1, + }, + } + ) +) + + DEVICE_ARN = "baz" @@ -873,11 +972,15 @@ def test_bad_statistics(): dev = _aws_device(wires=1, foo="bar") tape = qml.tape.QuantumTape(measurements=[qml.classical_shadow(wires=[0])]) with pytest.raises(QuantumFunctionError, match="Unsupported return type:"): - dev.statistics(None, tape.measurements) + dev._statistics(None, tape.measurements) -def test_batch_execute_non_parallel(monkeypatch): +@patch.object(AwsDevice, "properties", new_callable=mock.PropertyMock) +def test_batch_execute_non_parallel(mock_properties, monkeypatch): """Test if the batch_execute() method simply calls the inherited method if parallel=False""" + mock_action = Mock() + mock_action.action = {"braket.ir.openqasm.program": None} + mock_properties.return_value = mock_action dev = _aws_device(wires=2, foo="bar", parallel=False) assert dev.parallel is False @@ -887,10 +990,14 @@ def test_batch_execute_non_parallel(monkeypatch): assert res == 1967 +@patch.object(AwsDevice, "properties", new_callable=mock.PropertyMock) @patch.object(AwsDevice, "run") -def test_batch_execute_non_parallel_tracker(mock_run): +def test_batch_execute_non_parallel_tracker(mock_run, mock_properties): """Tests tracking for a non-parallel batch""" mock_run.return_value = TASK + mock_action = Mock() + mock_action.action = {"braket.ir.openqasm.program": None} + mock_properties.return_value = mock_action dev = _aws_device(wires=2, foo="bar", parallel=False) with QuantumTape() as circuit: @@ -918,11 +1025,103 @@ def test_batch_execute_non_parallel_tracker(mock_run): callback.assert_called_with(latest=latest, history=history, totals=totals) +@patch.object(AwsDevice, "run") +def test_batch_execute_program_set(mock_run): + """Test batch_execute correctly runs program sets when they are supported""" + task = Mock() + task.result.return_value = PROGRAM_SET_RESULT + mock_run.return_value = task + dev = _aws_device(wires=4, foo="bar", parallel=False, supports_program_sets=True) + with QuantumTape() as circuit: + qml.Hadamard(wires=0) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliX(0) @ qml.PauliY(1)) + + circuits = [circuit, circuit] + result = dev.batch_execute(circuits) + + braket_circuit = Circuit().h(0).cnot(0, 1).i(2).i(3).ry(0, -anp.pi / 2).rx(1, anp.pi / 2) + mock_run.assert_called_with( + ProgramSet([braket_circuit, braket_circuit]), + s3_destination_folder=("foo", "bar"), + shots=SHOTS * 2, + poll_timeout_seconds=AwsQuantumTask.DEFAULT_RESULTS_POLL_TIMEOUT, + poll_interval_seconds=AwsQuantumTask.DEFAULT_RESULTS_POLL_INTERVAL, + foo="bar", + ) + assert len(result) == 2 + assert result[0] == 0.6 + assert result[1] == 0.6 + + +@patch.object(AwsDevice, "run") +def test_batch_execute_program_set_parametrize_differentiable(mock_run): + """Test batch_execute correctly runs program sets with trainable parameters""" + task = Mock() + task.result.return_value = PROGRAM_SET_RESULT + mock_run.return_value = task + dev = _aws_device( + wires=4, + foo="bar", + parallel=False, + parametrize_differentiable=True, + supports_program_sets=True, + ) + + with QuantumTape() as circuit1: + qml.Hadamard(wires=0) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliX(0) @ qml.PauliY(1)) + + with QuantumTape() as circuit2: + qml.Hadamard(wires=0) + qml.RX(0.123, wires=0) + qml.CNOT(wires=[0, 1]) + qml.sample(qml.PauliZ(0)) + + circuits = [circuit1, circuit2] + result = dev.batch_execute(circuits) + + braket_circuit1 = Circuit().h(0).cnot(0, 1).i(2).i(3).ry(0, -anp.pi / 2).rx(1, anp.pi / 2) + braket_circuit2 = Circuit().h(0).rx(0, FreeParameter("p_0")).cnot(0, 1).i(2).i(3) + mock_run.assert_called_with( + ProgramSet.zip([braket_circuit1, braket_circuit2], input_sets=[{}, {"p_0": 0.123}]), + s3_destination_folder=("foo", "bar"), + shots=SHOTS * 2, + poll_timeout_seconds=AwsQuantumTask.DEFAULT_RESULTS_POLL_TIMEOUT, + poll_interval_seconds=AwsQuantumTask.DEFAULT_RESULTS_POLL_INTERVAL, + foo="bar", + ) + assert len(result) == 2 + assert result[0] == 0.6 + assert ( + result[1] == [1, 1, -1, 1, -1, 1, -1, -1, -1, 1, -1, 1, 1, -1, -1, -1, -1, 1, -1, 1] + ).all() + + +def test_batch_execute_program_set_noncommuting(): + """Test batch_execute correctly runs program sets when they are supported""" + dev = _aws_device(wires=4, foo="bar", parallel=False, supports_program_sets=True) + with QuantumTape() as circuit: + qml.Hadamard(wires=0) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliX(0)) + qml.sample(qml.PauliY(0)) + + circuits = [circuit, circuit] + with pytest.raises(ValueError): + dev.batch_execute(circuits) + + +@patch.object(AwsDevice, "properties", new_callable=mock.PropertyMock) @patch.object(AwsDevice, "run_batch") -def test_aws_device_batch_execute_parallel(mock_run_batch): +def test_aws_device_batch_execute_parallel(mock_run_batch, mock_properties): """Test batch_execute(parallel=True) correctly calls batch execution methods for AwsDevices in Braket SDK""" mock_run_batch.return_value = TASK_BATCH + mock_action = Mock() + mock_action.action = {"braket.ir.openqasm.program": None} + mock_properties.return_value = mock_action dev = _aws_device(wires=4, foo="bar", parallel=True) assert dev.parallel is True @@ -1013,11 +1212,15 @@ def test_local_sim_batch_execute_parallel(mock_run_batch): ) +@patch.object(AwsDevice, "properties", new_callable=mock.PropertyMock) @patch.object(AwsDevice, "run_batch") -def test_aws_device_batch_execute_parallel_tracker(mock_run_batch): +def test_aws_device_batch_execute_parallel_tracker(mock_run_batch, mock_properties): """Asserts tracker updates during parallel execution for AWS devices""" mock_run_batch.return_value = TASK_BATCH + mock_action = Mock() + mock_action.action = {"braket.ir.openqasm.program": None} + mock_properties.return_value = mock_action type(TASK_BATCH).unsuccessful = PropertyMock(return_value={}) dev = _aws_device(wires=1, foo="bar", parallel=True) @@ -1081,8 +1284,9 @@ def test_local_sim_batch_execute_parallel_tracker(mock_run_batch): callback.assert_called_with(latest=latest, history=history, totals=totals) +@patch.object(AwsDevice, "properties", new_callable=mock.PropertyMock) @patch.object(AwsDevice, "run_batch") -def test_batch_execute_partial_fail_parallel_tracker(mock_run_batch): +def test_batch_execute_partial_fail_parallel_tracker(mock_run_batch, mock_properties): """Asserts tracker updates during a partial failure of parallel execution""" FAIL_TASK = Mock() @@ -1095,6 +1299,9 @@ def test_batch_execute_partial_fail_parallel_tracker(mock_run_batch): type(FAIL_BATCH).unsuccessful = PropertyMock(return_value={"failed_task_arn"}) mock_run_batch.return_value = FAIL_BATCH + mock_action = Mock() + mock_action.action = {"braket.ir.openqasm.program": None} + mock_properties.return_value = mock_action dev = _aws_device(wires=1, foo="bar", parallel=True) with QuantumTape() as circuit: @@ -1135,10 +1342,14 @@ def test_batch_execute_partial_fail_parallel_tracker(mock_run_batch): callback.assert_called_with(latest=latest, history=history, totals=totals) +@patch.object(AwsDevice, "properties", new_callable=mock.PropertyMock) @patch.object(AwsDevice, "run_batch") -def test_batch_execute_parametrize_differentiable(mock_run_batch): +def test_batch_execute_parametrize_differentiable(mock_run_batch, mock_properties): """Test batch_execute(parallel=True) correctly calls batch execution methods in Braket SDK""" mock_run_batch.return_value = TASK_BATCH + mock_action = Mock() + mock_action.action = {"braket.ir.openqasm.program": None} + mock_properties.return_value = mock_action dev = _aws_device(wires=4, foo="bar", parametrize_differentiable=True, parallel=True) with QuantumTape() as circuit1: @@ -1690,6 +1901,20 @@ def test_run_task_unimplemented(): dev.execute(circuit) +def test_run_batch_task_unimplemented(): + """Tests that an error is thrown when _run_task_batch is not implemented""" + dummy = DummyCircuitSimulator() + dev = DummyLocalQubitDevice(wires=2, device=dummy, shots=1000, parallel=True) + + with QuantumTape() as circuit: + qml.Hadamard(wires=0) + qml.CNOT(wires=[0, 1]) + qml.probs(wires=[0, 1]) + + with pytest.raises(NotImplementedError): + dev.batch_execute([circuit, circuit]) + + @patch("braket.pennylane_plugin.braket_device.AwsDevice") def test_add_braket_user_agent_invoked(aws_device_mock): aws_device_mock_instance = aws_device_mock.return_value @@ -2020,12 +2245,12 @@ def _aws_device( action_properties=ACTION_PROPERTIES, native_gate_set=None, parametrize_differentiable=False, + supports_program_sets=False, **kwargs, ): properties_mock.action = {DeviceActionType.OPENQASM: action_properties} - properties_mock.return_value.action.return_value = { - DeviceActionType.OPENQASM: action_properties - } + if supports_program_sets: + properties_mock.action[DeviceActionType.OPENQASM_PROGRAM_SET] = action_properties properties_mock.paradigm.nativeGateSet = native_gate_set if native_gate_set is None: type(properties_mock).paradigm = PropertyMock(side_effect=AttributeError("paradigm")) @@ -2445,9 +2670,11 @@ def test_multiple_simultaneous_pulses_on_a_wire_raises_error(self): @patch.object(AwsDevice, "run_batch") @patch.object(AwsDevice, "name", new_callable=mock.PropertyMock) +@patch.object(AwsDevice, "properties", new_callable=mock.PropertyMock) @patch.object(BraketAwsQubitDevice, "_braket_to_pl_result") def test_batch_execute_with_noise_model( mock_to_result, + mock_properties, mock_name, mock_run_batch, noise_model, @@ -2456,6 +2683,9 @@ def test_batch_execute_with_noise_model( ): NUM_CIRCUITS = 5 mock_name.return_value = "dm1" + mock_action = Mock() + mock_action.action = {"braket.ir.openqasm.program": None} + mock_properties.return_value = mock_action dev = _aws_device( wires=4, device_type=AwsDeviceType.SIMULATOR, @@ -2492,8 +2722,12 @@ def test_verbatim_unsupported(device_type): @pytest.mark.parametrize("device_type", (AwsDeviceType.QPU, AwsDeviceType.SIMULATOR)) +@patch.object(AwsDevice, "properties", new_callable=mock.PropertyMock) @patch.object(AwsDevice, "run") -def test_native(mock_run, device_type): +def test_native(mock_run, mock_properties, device_type): + mock_action = Mock() + mock_action.action = {"braket.ir.openqasm.program": None} + mock_properties.return_value = mock_action dev = _aws_device( wires=2, device_arn="foo", diff --git a/test/unit_tests/test_shadow_expval.py b/test/unit_tests/test_shadow_expval.py index 4dcc6237..95ba453b 100644 --- a/test/unit_tests/test_shadow_expval.py +++ b/test/unit_tests/test_shadow_expval.py @@ -14,9 +14,10 @@ ) from braket.device_schema.simulators import GateModelSimulatorDeviceCapabilities from braket.devices import LocalSimulator +from braket.program_sets import ProgramSet from braket.simulator import BraketSimulator -from braket.task_result import GateModelTaskResult -from braket.tasks import GateModelQuantumTaskResult +from braket.task_result import GateModelTaskResult, ProgramSetTaskResult +from braket.tasks import GateModelQuantumTaskResult, ProgramSetQuantumTaskResult from pennylane.measurements import MeasurementTransform from pennylane.tape import QuantumScript, QuantumTape from pennylane.wires import Wires @@ -173,6 +174,114 @@ } ) ) +PROGRAM_RESULT_1 = { + "braketSchemaHeader": { + "name": "braket.task_result.program_result", + "version": "1", + }, + "executableResults": [ + { + "braketSchemaHeader": { + "name": "braket.task_result.program_set_executable_result", + "version": "1", + }, + "measurements": [SNAPSHOTS[0]], + "measuredQubits": [0, 1], + "inputsIndex": 0, + } + ], + "source": { + "braketSchemaHeader": { + "name": "braket.ir.openqasm.program", + "version": "1", + }, + "source": "OPENQASM 3.0;\nbit[2] b;\nqubit[2] q;\nh q[0];\ncnot q[0], q[1];\nb[0] = measure q[0];\nb[1] = measure q[1];", # noqa + "inputs": {"theta": [0.12, 2.1]}, + }, + "additionalMetadata": { + "simulatorMetadata": { + "braketSchemaHeader": { + "name": "braket.task_result.simulator_metadata", + "version": "1", + }, + "executionDuration": 50, + } + }, +} +PROGRAM_RESULT_2 = { + "braketSchemaHeader": { + "name": "braket.task_result.program_result", + "version": "1", + }, + "executableResults": [ + { + "braketSchemaHeader": { + "name": "braket.task_result.program_set_executable_result", + "version": "1", + }, + "measurements": [SNAPSHOTS[1]], + "measuredQubits": [0, 1], + "inputsIndex": 0, + } + ], + "source": { + "braketSchemaHeader": { + "name": "braket.ir.openqasm.program", + "version": "1", + }, + "source": "OPENQASM 3.0;\nbit[2] b;\nqubit[2] q;\nh q[0];\ncnot q[0], q[1];\nb[0] = measure q[0];\nb[1] = measure q[1];", # noqa + "inputs": {"theta": [0.12, 2.1]}, + }, + "additionalMetadata": { + "simulatorMetadata": { + "braketSchemaHeader": { + "name": "braket.task_result.simulator_metadata", + "version": "1", + }, + "executionDuration": 50, + } + }, +} +PROGRAM_SET_RESULT = ProgramSetQuantumTaskResult.from_object( + ProgramSetTaskResult( + **{ + "braketSchemaHeader": { + "name": "braket.task_result.program_set_task_result", + "version": "1", + }, + "programResults": [PROGRAM_RESULT_1, PROGRAM_RESULT_2], + "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": 2, + "successfulShots": 2, + "programMetadata": [{"executables": [{}]}], + "deviceParameters": { + "braketSchemaHeader": { + "name": "braket.device_schema.simulators.gate_model_simulator_device_parameters", + "version": "1", + }, + "paradigmParameters": { + "braketSchemaHeader": { + "name": "braket.device_schema.gate_model_parameters", + "version": "1", + }, + "qubitCount": 2, + "disableQubitRewiring": False, + }, + }, + "createdAt": "2024-10-15T19:06:58.986Z", + "endedAt": "2024-10-15T19:07:00.382Z", + "status": "COMPLETED", + "totalFailedExecutables": 0, + }, + } + ) +) TASK = Mock() TASK.result.return_value = RESULT type(TASK).id = PropertyMock(return_value="task_arn") @@ -180,6 +289,8 @@ TASK_BATCH = Mock() TASK_BATCH.results.return_value = [BATCH_RESULT_1, BATCH_RESULT_2] type(TASK_BATCH).tasks = PropertyMock(return_value=[TASK, TASK]) +TASK_PROGRAM_SET = Mock() +TASK_PROGRAM_SET.result.return_value = PROGRAM_SET_RESULT @pytest.mark.xfail(raises=ValueError) @@ -212,58 +323,52 @@ def test_only_one_operator_in_shadow_expval(): circs[0].h(1) +@patch.object(AwsDevice, "properties", new_callable=mock.PropertyMock) @patch.object(AwsDevice, "run") @patch.object(AwsDevice, "run_batch") @pytest.mark.parametrize( - "pl_circ, wires, result_types, expected_pl_result", + "pl_circ, wires, expected_pl_result", [ ( CIRCUIT_1, 2, - [ - { - "type": { - "observable": "x()", - "targets": [[0]], - "parameters": ["p_0", "p_1"], - "type": "expectation_value", - }, - "value": { - "expectation": 1.5, - }, - }, - ], [[1.5]], ), ], ) @pytest.mark.parametrize( - "parallel, expected_braket_circ, return_val, call_count, max_para, max_conn", + "parallel, expected_braket_task_spec, return_val, call_count, max_para, max_conn, supports_program_sets", [ - (False, circs, TASK, SHOTS, None, None), - (True, circs, TASK_BATCH, 1, 10, 10), + (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), ], ) def test_shadow_expval_aws_device( mock_run_batch, mock_run, + mock_properties, pl_circ, wires, - result_types, expected_pl_result, parallel, - expected_braket_circ, + expected_braket_task_spec, return_val, call_count, max_para, max_conn, + supports_program_sets, ): + mock_action = Mock() + mock_action.action = {"braket.ir.openqasm.program": None} + mock_properties.return_value = mock_action dev = _aws_device( wires=2, foo="bar", parallel=parallel, max_parallel=max_para, max_connections=max_conn, + supports_program_sets=supports_program_sets, ) mock_runner = mock_run_batch if parallel else mock_run mock_runner.return_value = return_val @@ -274,7 +379,7 @@ def test_shadow_expval_aws_device( assert res == expected_pl_result[0] if parallel: mock_runner.assert_called_with( - expected_braket_circ, + expected_braket_task_spec, s3_destination_folder=("foo", "bar"), shots=1, **kwargs, @@ -282,8 +387,18 @@ def test_shadow_expval_aws_device( poll_interval_seconds=AwsQuantumTask.DEFAULT_RESULTS_POLL_INTERVAL, foo="bar", ) + elif supports_program_sets: + mock_runner.assert_called_with( + expected_braket_task_spec, + s3_destination_folder=("foo", "bar"), + shots=2, + **kwargs, + poll_timeout_seconds=AwsQuantumTask.DEFAULT_RESULTS_POLL_TIMEOUT, + poll_interval_seconds=AwsQuantumTask.DEFAULT_RESULTS_POLL_INTERVAL, + foo="bar", + ) else: - for c in expected_braket_circ: + for c in expected_braket_task_spec: mock_runner.assert_any_call( c, s3_destination_folder=("foo", "bar"), @@ -297,25 +412,12 @@ def test_shadow_expval_aws_device( @patch.object(LocalSimulator, "run") @pytest.mark.parametrize( - "pl_circ, expected_braket_circ, wires, result_types, expected_pl_result", + "pl_circ, expected_braket_circ, wires, expected_pl_result", [ ( CIRCUIT_1, Circuit().h(0).cnot(0, 1).rx(0, 0.432).ry(0, 0.543), 2, - [ - { - "type": { - "observable": "x()", - "targets": [[0]], - "parameters": ["p_0", "p_1"], - "type": "expectation_value", - }, - "value": { - "expectation": 1.5, - }, - }, - ], [[1.5]], ), ], @@ -326,7 +428,6 @@ def test_shadow_expval_local( pl_circ, expected_braket_circ, wires, - result_types, expected_pl_result, backend, ): @@ -360,12 +461,12 @@ def _aws_device( shots=SHOTS, device_arn="baz", action_properties=ACTION_PROPERTIES, + supports_program_sets=False, **kwargs, ): properties_mock.action = {DeviceActionType.OPENQASM: action_properties} - properties_mock.return_value.action.return_value = { - DeviceActionType.OPENQASM: action_properties - } + if supports_program_sets: + properties_mock.action[DeviceActionType.OPENQASM_PROGRAM_SET] = action_properties type_mock.return_value = device_type dev = BraketAwsQubitDevice( wires=wires, @@ -468,12 +569,16 @@ def test_run_snapshots_not_implemented(): dev.execute(circuit) +@patch.object(AwsDevice, "properties", new_callable=mock.PropertyMock) @patch.object(AwsDevice, "run_batch") -def test_shadows_parallel_tracker(mock_run_batch): +def test_shadows_parallel_tracker(mock_run_batch, mock_properties): """Asserts tracker updates during parallel shadows computation""" mock_run_batch.return_value = TASK_BATCH type(TASK_BATCH).unsuccessful = PropertyMock(return_value={}) + mock_action = Mock() + mock_action.action = {"braket.ir.openqasm.program": None} + mock_properties.return_value = mock_action dev = _aws_device(wires=2, foo="bar", parallel=True, shots=SHOTS) with QuantumTape() as circuit: