diff --git a/src/braket/pennylane_plugin/braket_device.py b/src/braket/pennylane_plugin/braket_device.py index 06ce9748..e0abd675 100644 --- a/src/braket/pennylane_plugin/braket_device.py +++ b/src/braket/pennylane_plugin/braket_device.py @@ -48,6 +48,7 @@ from pennylane.exceptions import QuantumFunctionError from pennylane.gradients import param_shift from pennylane.measurements import ( + ClassicalShadowMP, CountsMP, ExpectationMP, MeasurementProcess, @@ -172,6 +173,11 @@ def __init__( DeviceActionType.OPENQASM_PROGRAM_SET in self._device.properties.action and self._shots is not None ) + self._max_program_set_executables = ( + self._device.properties.action["braket.ir.openqasm.program_set"].maximumExecutables + if self._supports_program_sets + else None + ) if noise_model: self._validate_noise_model_support() @@ -206,13 +212,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: - return super().batch_execute(circuits) - - if self._supports_program_sets and ( - len(circuits) - > self._device.properties.action["braket.ir.openqasm.program_set"].maximumExecutables - ): + if not self._parallel and not self._can_run_as_program_set(len(circuits)): return super().batch_execute(circuits) for circuit in circuits: @@ -220,6 +220,17 @@ def batch_execute(self, circuits, **run_kwargs): all_trainable = [] braket_circuits = [] for circuit in circuits: + if isinstance(circuit.observables[0], ClassicalShadowMP): + if len(circuit.observables) > 1: + raise ValueError( + "A circuit with a ClassicalShadowMP observable must " + "have that as its only return type." + ) + if len(circuits) > 1: + raise ValueError("Classical shadow must be called with a single circuit.") + bits, recipes = self.classical_shadow(circuit.observables[0], circuit) + return [(bits, recipes)] + trainable = ( BraketQubitDevice._get_trainable_parameters(circuit) if self._parametrize_differentiable @@ -230,7 +241,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=not self._can_run_as_program_set(len(circuits)), **run_kwargs, ) ) @@ -245,6 +256,10 @@ def batch_execute(self, circuits, **run_kwargs): return self._run_task_batch(braket_circuits, circuits, batch_shots, batch_inputs) + def _can_run_as_program_set(self, n_executables: int) -> bool: + """Whether a list of executables can be run as a program set""" + return self._supports_program_sets and n_executables < self._max_program_set_executables + def _pl_to_braket_circuit( self, circuit: QuantumTape, @@ -510,7 +525,7 @@ def execute(self, circuit: QuantumTape, compute_gradient=False, **run_kwargs) -> if len(circuit.observables) > 1: raise ValueError( "A circuit with a ShadowExpvalMP observable must " - "have that as its only result type." + "have that as its only return type." ) return [self.shadow_expval(circuit.observables[0], circuit)] raise RuntimeError("The circuit has an unsupported MeasurementTransform.") @@ -707,7 +722,7 @@ def use_grouping(self) -> bool: return not (caps.get("provides_jacobian")) def _run_task_batch(self, braket_circuits, pl_circuits, batch_shots: int, inputs): - if self._supports_program_sets: + if self._can_run_as_program_set(len(braket_circuits)): program_set = ( ProgramSet.zip(braket_circuits, input_sets=inputs) if inputs @@ -764,7 +779,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._can_run_as_program_set(n_snapshots): program_set = ProgramSet(snapshot_circuits) task = self._device.run( program_set, diff --git a/test/unit_tests/test_braket_device.py b/test/unit_tests/test_braket_device.py index ec7475f1..3bd7cb19 100644 --- a/test/unit_tests/test_braket_device.py +++ b/test/unit_tests/test_braket_device.py @@ -1125,8 +1125,8 @@ 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): +@patch.object(AwsDevice, "run_batch") +def test_batch_execute_program_set_exceeds_max_executables(mock_run_batch): """Test batch_execute falls back to individual programs when exceeding maximumExecutables""" custom_result = GateModelQuantumTaskResult.from_string( json.dumps( @@ -1174,7 +1174,12 @@ def test_batch_execute_program_set_exceeds_max_executables(mock_run): task.result.return_value = custom_result type(task).id = PropertyMock(return_value="task_arn") task.state.return_value = "COMPLETED" - mock_run.return_value = task + + # Mock the batch to return 101 results (one for each circuit) + task_batch = Mock() + task_batch.results.return_value = [custom_result] * 101 + type(task_batch).tasks = PropertyMock(return_value=[task] * 101) + mock_run_batch.return_value = task_batch dev = _aws_device(wires=4, foo="bar", parallel=True, supports_program_sets=True) @@ -1199,7 +1204,9 @@ def test_batch_execute_program_set_exceeds_max_executables(mock_run): ) result = dev.batch_execute(circuits) - assert mock_run.call_count == 101 + assert mock_run_batch.call_count == 1 + call_args = mock_run_batch.call_args[0] + assert len(call_args[0]) == 101 # First argument is the list of circuits assert len(result) == 101 diff --git a/test/unit_tests/test_shadow_expval.py b/test/unit_tests/test_shadow_expval.py index 84a395bd..277bc040 100644 --- a/test/unit_tests/test_shadow_expval.py +++ b/test/unit_tests/test_shadow_expval.py @@ -22,8 +22,9 @@ from braket.aws import AwsDevice, AwsDeviceType, AwsQuantumTask from braket.circuits import Circuit from braket.device_schema import DeviceActionType -from braket.device_schema.openqasm_device_action_properties import ( +from braket.device_schema import ( OpenQASMDeviceActionProperties, + OpenQASMProgramSetDeviceActionProperties, ) from braket.device_schema.simulators import GateModelSimulatorDeviceCapabilities from braket.devices import LocalSimulator @@ -66,6 +67,17 @@ ) ) +ACTION_PROPERTIES_PROGRAM_SET = OpenQASMProgramSetDeviceActionProperties.parse_raw( + json.dumps( + { + "version": ["1"], + "actionType": "braket.ir.openqasm.program_set", + "maximumExecutables": 10, + "maximumTotalShots": 200000, + } + ) +) + SNAPSHOTS = [[0, 0], [1, 1]] GATE_MODEL_RESULT = GateModelTaskResult( @@ -335,6 +347,12 @@ def test_only_one_operator_in_shadow_expval(): # basis rotation with seed circs[0].h(1) +# Circuits with measurements for program set (required for program sets to work) +circs_with_measurements = [ + Circuit().h(0).cnot(0, 1).rx(0, 0.432).ry(0, 0.543).h(1).measure(0).measure(1), + Circuit().h(0).cnot(0, 1).rx(0, 0.432).ry(0, 0.543).measure(0).measure(1), +] + @patch.object(AwsDevice, "properties", new_callable=mock.PropertyMock) @patch.object(AwsDevice, "run") @@ -373,7 +391,13 @@ def test_shadow_expval_aws_device( supports_program_sets, ): mock_action = Mock() - mock_action.action = {"braket.ir.openqasm.program": None} + if supports_program_sets: + mock_action.action = { + DeviceActionType.OPENQASM: ACTION_PROPERTIES, + DeviceActionType.OPENQASM_PROGRAM_SET: ACTION_PROPERTIES_PROGRAM_SET, + } + else: + mock_action.action = {DeviceActionType.OPENQASM: ACTION_PROPERTIES} mock_properties.return_value = mock_action dev = _aws_device( wires=2, @@ -479,7 +503,9 @@ def _aws_device( ): properties_mock.action = {DeviceActionType.OPENQASM: action_properties} if supports_program_sets: - properties_mock.action[DeviceActionType.OPENQASM_PROGRAM_SET] = action_properties + properties_mock.action[DeviceActionType.OPENQASM_PROGRAM_SET] = ( + ACTION_PROPERTIES_PROGRAM_SET + ) type_mock.return_value = device_type dev = BraketAwsQubitDevice( wires=wires, @@ -649,3 +675,82 @@ def test_non_shadow_expval_transform(): dummy_measurement_transform() dev.execute(circuit) + + +# Test for classical_shadow with program sets +CIRCUIT_CLASSICAL_SHADOW = QuantumScript( + ops=[ + qml.Hadamard(wires=0), + qml.CNOT(wires=[0, 1]), + ], + measurements=[qml.classical_shadow(wires=range(2), seed=SEED)], +) + + +@patch.object(AwsDevice, "properties", new_callable=mock.PropertyMock) +@patch.object(AwsDevice, "run") +def test_batch_execute_classical_shadow_single_circuit(mock_run, mock_properties): + """Test that batch_execute handles ClassicalShadowMP with a single circuit""" + mock_action = Mock() + mock_action.action = { + DeviceActionType.OPENQASM: ACTION_PROPERTIES, + DeviceActionType.OPENQASM_PROGRAM_SET: ACTION_PROPERTIES_PROGRAM_SET, + } + mock_properties.return_value = mock_action + + dev = _aws_device( + wires=2, + foo="bar", + supports_program_sets=True, + ) + + mock_run.return_value = TASK_PROGRAM_SET + circuits = [CIRCUIT_CLASSICAL_SHADOW] + results = dev.batch_execute(circuits) + + assert len(results) == 1 + bits, recipes = results[0] + assert bits.shape == (SHOTS, 2) + assert recipes.shape == (SHOTS, 2) + + +@patch.object(AwsDevice, "properties", new_callable=mock.PropertyMock) +@patch.object(AwsDevice, "run") +@pytest.mark.parametrize( + "circuits, error_message", + [ + ( + [ + QuantumScript( + ops=[qml.Hadamard(wires=0)], + measurements=[ + qml.classical_shadow(wires=[0], seed=SEED), + qml.expval(qml.PauliZ(0)), + ], + ) + ], + "A circuit with a ClassicalShadowMP observable must have that as its only return type", + ), + ( + [CIRCUIT_CLASSICAL_SHADOW, CIRCUIT_CLASSICAL_SHADOW], + "Classical shadow must be called with a single circuit", + ), + ], +) +def test_batch_execute_classical_shadow_errors(mock_run, mock_properties, circuits, error_message): + """Test that batch_execute raises appropriate errors for invalid ClassicalShadowMP usage""" + mock_action = Mock() + mock_action.action = { + DeviceActionType.OPENQASM: ACTION_PROPERTIES, + DeviceActionType.OPENQASM_PROGRAM_SET: ACTION_PROPERTIES_PROGRAM_SET, + } + mock_properties.return_value = mock_action + + dev = _aws_device( + wires=2, + foo="bar", + supports_program_sets=True, + ) + + with pytest.raises(ValueError, match=error_message): + dev.batch_execute(circuits)