From 74081e1611feb7b4c7c8b45159bf4d1d77620b35 Mon Sep 17 00:00:00 2001 From: Tim Chen Date: Thu, 13 Nov 2025 19:06:24 -0500 Subject: [PATCH 1/4] fix: classical shadow with program set --- src/braket/pennylane_plugin/braket_device.py | 12 ++ test/unit_tests/test_shadow_expval.py | 144 +++++++++++++++++++ 2 files changed, 156 insertions(+) diff --git a/src/braket/pennylane_plugin/braket_device.py b/src/braket/pennylane_plugin/braket_device.py index 159e7123..2a17297d 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, @@ -220,6 +221,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 result 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 diff --git a/test/unit_tests/test_shadow_expval.py b/test/unit_tests/test_shadow_expval.py index 84a395bd..5c74d4af 100644 --- a/test/unit_tests/test_shadow_expval.py +++ b/test/unit_tests/test_shadow_expval.py @@ -335,6 +335,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") @@ -649,3 +655,141 @@ 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_classical_shadow_with_program_sets(mock_run, mock_properties): + """Test that classical_shadow measurement raises error (not yet supported in execute)""" + mock_action = Mock() + mock_action.action = {"braket.ir.openqasm.program": None} + mock_properties.return_value = mock_action + + dev = _aws_device( + wires=2, + foo="bar", + supports_program_sets=True, + ) + + # When using program sets, return TASK_PROGRAM_SET instead of TASK + mock_run.return_value = TASK_PROGRAM_SET + + # Execute circuit with classical_shadow should raise an error + # because ClassicalShadowMP is not yet handled in the execute method + with pytest.raises( + RuntimeError, + match="The circuit has an unsupported MeasurementTransform", + ): + dev.execute(CIRCUIT_CLASSICAL_SHADOW) + + +@patch.object(AwsDevice, "properties", new_callable=mock.PropertyMock) +@patch.object(AwsDevice, "run") +def test_classical_shadow_multiple_observables_error(mock_run, mock_properties): + """Test that classical_shadow raises error (not yet supported in execute)""" + mock_action = Mock() + mock_action.action = {"braket.ir.openqasm.program": None} + mock_properties.return_value = mock_action + + dev = _aws_device(wires=2, foo="bar") + + # Create circuit with multiple measurements including classical_shadow + circuit = QuantumScript( + ops=[qml.Hadamard(wires=0)], + measurements=[ + qml.classical_shadow(wires=[0], seed=SEED), + qml.expval(qml.PauliZ(0)), + ], + ) + + # Since ClassicalShadowMP is not handled in execute, it raises RuntimeError + with pytest.raises( + RuntimeError, + match="The circuit has an unsupported MeasurementTransform", + ): + dev.execute(circuit) + + +@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_program_set_action = Mock() + mock_program_set_action.maximumExecutables = 10 + + mock_action = Mock() + mock_action.action = { + "braket.ir.openqasm.program": None, + "braket.ir.openqasm.program_set": mock_program_set_action, + } + 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 result 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_program_set_action = Mock() + mock_program_set_action.maximumExecutables = 10 + + mock_action = Mock() + mock_action.action = { + "braket.ir.openqasm.program": None, + "braket.ir.openqasm.program_set": mock_program_set_action, + } + 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) From 73f236a1d1cbeb05e8e54fe1dc9714545d00341a Mon Sep 17 00:00:00 2001 From: Tim Chen Date: Fri, 14 Nov 2025 09:59:42 -0500 Subject: [PATCH 2/4] fallback to programs beyond max exectuables --- src/braket/pennylane_plugin/braket_device.py | 12 +++--- test/unit_tests/test_shadow_expval.py | 45 ++++++++++++++------ 2 files changed, 38 insertions(+), 19 deletions(-) diff --git a/src/braket/pennylane_plugin/braket_device.py b/src/braket/pennylane_plugin/braket_device.py index 2a17297d..cd2dfe3e 100644 --- a/src/braket/pennylane_plugin/braket_device.py +++ b/src/braket/pennylane_plugin/braket_device.py @@ -173,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() @@ -210,10 +215,7 @@ 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 self._supports_program_sets and len(circuits) > self._max_program_set_executables: return super().batch_execute(circuits) for circuit in circuits: @@ -776,7 +778,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._supports_program_sets and n_snapshots < self._max_program_set_executables: program_set = ProgramSet(snapshot_circuits) task = self._device.run( program_set, diff --git a/test/unit_tests/test_shadow_expval.py b/test/unit_tests/test_shadow_expval.py index 5c74d4af..c91dcd4e 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( @@ -379,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, @@ -485,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, @@ -672,7 +692,10 @@ def test_non_shadow_expval_transform(): def test_classical_shadow_with_program_sets(mock_run, mock_properties): """Test that classical_shadow measurement raises error (not yet supported in execute)""" mock_action = Mock() - mock_action.action = {"braket.ir.openqasm.program": None} + mock_action.action = { + DeviceActionType.OPENQASM: ACTION_PROPERTIES, + DeviceActionType.OPENQASM_PROGRAM_SET: ACTION_PROPERTIES_PROGRAM_SET, + } mock_properties.return_value = mock_action dev = _aws_device( @@ -724,13 +747,10 @@ def test_classical_shadow_multiple_observables_error(mock_run, mock_properties): @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_program_set_action = Mock() - mock_program_set_action.maximumExecutables = 10 - mock_action = Mock() mock_action.action = { - "braket.ir.openqasm.program": None, - "braket.ir.openqasm.program_set": mock_program_set_action, + DeviceActionType.OPENQASM: ACTION_PROPERTIES, + DeviceActionType.OPENQASM_PROGRAM_SET: ACTION_PROPERTIES_PROGRAM_SET, } mock_properties.return_value = mock_action @@ -775,13 +795,10 @@ def test_batch_execute_classical_shadow_single_circuit(mock_run, mock_properties ) 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_program_set_action = Mock() - mock_program_set_action.maximumExecutables = 10 - mock_action = Mock() mock_action.action = { - "braket.ir.openqasm.program": None, - "braket.ir.openqasm.program_set": mock_program_set_action, + DeviceActionType.OPENQASM: ACTION_PROPERTIES, + DeviceActionType.OPENQASM_PROGRAM_SET: ACTION_PROPERTIES_PROGRAM_SET, } mock_properties.return_value = mock_action From 7e1b2478cae3389d047d5f6f7a55f739cd82b45e Mon Sep 17 00:00:00 2001 From: Tim Chen Date: Fri, 14 Nov 2025 10:13:45 -0500 Subject: [PATCH 3/4] remove outdated tests --- test/unit_tests/test_shadow_expval.py | 56 --------------------------- 1 file changed, 56 deletions(-) diff --git a/test/unit_tests/test_shadow_expval.py b/test/unit_tests/test_shadow_expval.py index c91dcd4e..fe47d1dc 100644 --- a/test/unit_tests/test_shadow_expval.py +++ b/test/unit_tests/test_shadow_expval.py @@ -687,62 +687,6 @@ def test_non_shadow_expval_transform(): ) -@patch.object(AwsDevice, "properties", new_callable=mock.PropertyMock) -@patch.object(AwsDevice, "run") -def test_classical_shadow_with_program_sets(mock_run, mock_properties): - """Test that classical_shadow measurement raises error (not yet supported in execute)""" - 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, - ) - - # When using program sets, return TASK_PROGRAM_SET instead of TASK - mock_run.return_value = TASK_PROGRAM_SET - - # Execute circuit with classical_shadow should raise an error - # because ClassicalShadowMP is not yet handled in the execute method - with pytest.raises( - RuntimeError, - match="The circuit has an unsupported MeasurementTransform", - ): - dev.execute(CIRCUIT_CLASSICAL_SHADOW) - - -@patch.object(AwsDevice, "properties", new_callable=mock.PropertyMock) -@patch.object(AwsDevice, "run") -def test_classical_shadow_multiple_observables_error(mock_run, mock_properties): - """Test that classical_shadow raises error (not yet supported in execute)""" - mock_action = Mock() - mock_action.action = {"braket.ir.openqasm.program": None} - mock_properties.return_value = mock_action - - dev = _aws_device(wires=2, foo="bar") - - # Create circuit with multiple measurements including classical_shadow - circuit = QuantumScript( - ops=[qml.Hadamard(wires=0)], - measurements=[ - qml.classical_shadow(wires=[0], seed=SEED), - qml.expval(qml.PauliZ(0)), - ], - ) - - # Since ClassicalShadowMP is not handled in execute, it raises RuntimeError - with pytest.raises( - RuntimeError, - match="The circuit has an unsupported MeasurementTransform", - ): - dev.execute(circuit) - - @patch.object(AwsDevice, "properties", new_callable=mock.PropertyMock) @patch.object(AwsDevice, "run") def test_batch_execute_classical_shadow_single_circuit(mock_run, mock_properties): From 001d4da2172023c3d00c507af772a32fcf772542 Mon Sep 17 00:00:00 2001 From: Tim Chen Date: Mon, 17 Nov 2025 10:42:51 -0500 Subject: [PATCH 4/4] address the comment --- src/braket/pennylane_plugin/braket_device.py | 19 ++++++++++--------- test/unit_tests/test_braket_device.py | 15 +++++++++++---- test/unit_tests/test_shadow_expval.py | 2 +- 3 files changed, 22 insertions(+), 14 deletions(-) diff --git a/src/braket/pennylane_plugin/braket_device.py b/src/braket/pennylane_plugin/braket_device.py index cd2dfe3e..a4041f02 100644 --- a/src/braket/pennylane_plugin/braket_device.py +++ b/src/braket/pennylane_plugin/braket_device.py @@ -212,10 +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._max_program_set_executables: + if not self._parallel and not self._can_run_as_program_set(len(circuits)): return super().batch_execute(circuits) for circuit in circuits: @@ -227,7 +224,7 @@ def batch_execute(self, circuits, **run_kwargs): if len(circuit.observables) > 1: raise ValueError( "A circuit with a ClassicalShadowMP observable must " - "have that as its only result type." + "have that as its only return type." ) if len(circuits) > 1: raise ValueError("Classical shadow must be called with a single circuit.") @@ -244,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, ) ) @@ -259,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, @@ -524,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.") @@ -721,7 +722,7 @@ def use_grouping(self) -> bool: return not ("provides_jacobian" in caps and caps["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 @@ -778,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 and n_snapshots < self._max_program_set_executables: + 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 aaac7a41..8fd4326a 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 fe47d1dc..277bc040 100644 --- a/test/unit_tests/test_shadow_expval.py +++ b/test/unit_tests/test_shadow_expval.py @@ -729,7 +729,7 @@ def test_batch_execute_classical_shadow_single_circuit(mock_run, mock_properties ], ) ], - "A circuit with a ClassicalShadowMP observable must have that as its only result type", + "A circuit with a ClassicalShadowMP observable must have that as its only return type", ), ( [CIRCUIT_CLASSICAL_SHADOW, CIRCUIT_CLASSICAL_SHADOW],