Error in classical_shadow as return type when using program sets on QPU #313
Description
Describe the bug
An error is encountered when using program sets (implicitly invoked) on QPU with the qml.classical_shadow(wires) as return type.
To reproduce
import pennylane as qml
from pennylane import numpy as np
from functools import partial
dev = qml.device(
"braket.aws.qubit",
device_arn="arn:aws:braket:us-west-1::device/qpu/rigetti/Ankaa-3",
wires=2,
)
@partial(qml.set_shots, shots=10)
@qml.qnode(dev)
def bell_circuit():
qml.Hadamard(0)
qml.CNOT((0,1))
return qml.classical_shadow(wires=range(2))
bits, recipes = bell_circuit()
Expected behavior
When dev = qml.device("default.qubit", wires=2) is set as the backend, the above code could run successfully without any error. One would expect the above code to run normally on a real QPU.
Screenshots or logs
---------------------------------------------------------------------------
AttributeError Traceback (most recent call last)
Cell In[7], line 16
13 qml.CNOT((0, 1))
14 return qml.classical_shadow(wires=range(2))
---> 16 bits, recipes = bell_circuit()
File ~/anaconda3/envs/Braket/lib/python3.10/site-packages/pennylane/workflow/qnode.py:922, in QNode.__call__(self, *args, **kwargs)
919 from ._capture_qnode import capture_qnode # pylint: disable=import-outside-toplevel
921 return capture_qnode(self, *args, **kwargs)
--> 922 return self._impl_call(*args, **kwargs)
File ~/anaconda3/envs/Braket/lib/python3.10/site-packages/pennylane/workflow/qnode.py:895, in QNode._impl_call(self, *args, **kwargs)
892 # Calculate the classical jacobians if necessary
893 self._transform_program.set_classical_component(self, args, kwargs)
--> 895 res = execute(
896 (tape,),
897 device=self.device,
898 diff_method=self.diff_method,
899 interface=self.interface,
900 transform_program=self._transform_program,
901 gradient_kwargs=self.gradient_kwargs,
902 **self.execute_kwargs,
903 )
904 res = res[0]
906 # convert result to the interface in case the qfunc has no parameters
File ~/anaconda3/envs/Braket/lib/python3.10/site-packages/pennylane/workflow/execution.py:233, in execute(tapes, device, diff_method, interface, grad_on_execution, cache, cachesize, max_diff, device_vjp, postselect_mode, mcm_method, gradient_kwargs, transform_program, executor_backend)
229 tapes, outer_post_processing = outer_transform(tapes)
231 assert not outer_transform.is_informative, "should only contain device preprocessing"
--> 233 results = run(tapes, device, config, inner_transform)
234 return user_post_processing(outer_post_processing(results))
File ~/anaconda3/envs/Braket/lib/python3.10/site-packages/pennylane/workflow/run.py:291, in run(tapes, device, config, inner_transform_program)
287 no_interface_boundary_required = (
288 config.interface == Interface.NUMPY or config.gradient_method == "backprop"
289 )
290 if no_interface_boundary_required:
--> 291 results = inner_execute(tapes)
292 return results
294 # TODO: Prune once support for tf-autograph is dropped
File ~/anaconda3/envs/Braket/lib/python3.10/site-packages/pennylane/workflow/run.py:256, in _make_inner_execute.<locals>.inner_execute(tapes)
253 transformed_tapes, transform_post_processing = inner_transform(tapes)
255 if transformed_tapes:
--> 256 results = device.execute(transformed_tapes, execution_config=execution_config)
257 else:
258 results = ()
File ~/anaconda3/envs/Braket/lib/python3.10/site-packages/pennylane/devices/modifiers/single_tape_support.py:30, in _make_execute.<locals>.execute(self, circuits, execution_config)
28 is_single_circuit = True
29 circuits = (circuits,)
---> 30 results = batch_execute(self, circuits, execution_config)
31 return results[0] if is_single_circuit else results
File ~/anaconda3/envs/Braket/lib/python3.10/site-packages/pennylane/devices/legacy_facade.py:378, in LegacyDeviceFacade.execute(self, circuits, execution_config)
376 first_shot = circuits[0].shots
377 if all(t.shots == first_shot for t in circuits):
--> 378 return _set_shots(dev, first_shot)(dev.batch_execute)(circuits, **kwargs)
379 return tuple(
380 _set_shots(dev, t.shots)(dev.batch_execute)((t,), **kwargs)[0] for t in circuits
381 )
File ~/anaconda3/envs/Braket/lib/python3.10/contextlib.py:79, in ContextDecorator.__call__.<locals>.inner(*args, **kwds)
76 @wraps(func)
77 def inner(*args, **kwds):
78 with self._recreate_cm():
---> 79 return func(*args, **kwds)
File ~/anaconda3/envs/Braket/lib/python3.10/site-packages/braket/pennylane_plugin/braket_device.py:240, in BraketQubitDevice.batch_execute(self, circuits, **run_kwargs)
232 batch_shots = 0 if self.analytic else self.shots
234 batch_inputs = (
235 [{f"p_{k}": v for k, v in trainable.items()} for trainable in all_trainable]
236 if self._parametrize_differentiable
237 else []
238 )
--> 240 return self._run_task_batch(braket_circuits, circuits, batch_shots, batch_inputs)
File ~/anaconda3/envs/Braket/lib/python3.10/site-packages/braket/pennylane_plugin/braket_device.py:718, in BraketAwsQubitDevice._run_task_batch(self, braket_circuits, pl_circuits, batch_shots, inputs)
705 program_set = (
706 ProgramSet.zip(braket_circuits, input_sets=inputs)
707 if inputs
708 else ProgramSet(braket_circuits)
709 )
710 task = self._device.run(
711 program_set,
712 s3_destination_folder=self._s3_folder,
(...)
716 **self._run_kwargs,
717 )
--> 718 return self._braket_program_set_to_pl_result(task.result(), pl_circuits)
719 task_batch = self._device.run_batch(
720 braket_circuits,
721 s3_destination_folder=self._s3_folder,
(...)
728 **self._run_kwargs,
729 )
731 # Call results() to retrieve the Braket results in parallel.
File ~/anaconda3/envs/Braket/lib/python3.10/site-packages/braket/pennylane_plugin/braket_device.py:409, in BraketQubitDevice._braket_program_set_to_pl_result(self, program_set_result, circuits)
405 measurements = program_result[0].measurements
407 # Program sets require shots > 0,
408 # so the circuit's measurements are guaranteed to be SampleMeasurements
--> 409 executable_results = [
410 measurement.process_samples(measurements, wire_order=measurement.wires)
411 for measurement in circuit.measurements
412 ]
413 results.append(
414 onp.array(executable_results).squeeze()
415 if len(circuit.measurements) == 1
416 else tuple(onp.array(result).squeeze() for result in executable_results)
417 )
418 return results
File ~/anaconda3/envs/Braket/lib/python3.10/site-packages/braket/pennylane_plugin/braket_device.py:410, in <listcomp>(.0)
405 measurements = program_result[0].measurements
407 # Program sets require shots > 0,
408 # so the circuit's measurements are guaranteed to be SampleMeasurements
409 executable_results = [
--> 410 measurement.process_samples(measurements, wire_order=measurement.wires)
411 for measurement in circuit.measurements
412 ]
413 results.append(
414 onp.array(executable_results).squeeze()
415 if len(circuit.measurements) == 1
416 else tuple(onp.array(result).squeeze() for result in executable_results)
417 )
418 return results
AttributeError: 'ClassicalShadowMP' object has no attribute 'process_samples'
System information
A description of your system. Please provide:
- Amazon Braket Python PennyLane Plugin version: 1.33.4
- Amazon Braket Python SDK version: 1.103.0
- Amazon Braket Python Schemas version: 1.26.1
- Amazon Braket Python Default Simulator version: 1.32.0
- Python version: 3.10.19
Additional context
Add any other context about the problem here.