fix: classical shadow with program set

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,20 +212,25 @@ 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:
             self.check_validity(circuit.operations, circuit.observables)
         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,

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
 
 

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)