From 4a6569d404e19476b980460432221bbbeb861d5d Mon Sep 17 00:00:00 2001 From: Cody Wang Date: Thu, 5 Sep 2024 19:20:01 -0700 Subject: [PATCH 1/6] feat: Support `SProd` and `Sum` for expval --- src/braket/pennylane_plugin/braket_device.py | 26 +++-- src/braket/pennylane_plugin/translation.py | 72 ++++++------ test/unit_tests/test_braket_device.py | 113 ++++++++----------- test/unit_tests/test_translation.py | 14 +-- 4 files changed, 106 insertions(+), 119 deletions(-) diff --git a/src/braket/pennylane_plugin/braket_device.py b/src/braket/pennylane_plugin/braket_device.py index 2fcc4a9e..48679b21 100644 --- a/src/braket/pennylane_plugin/braket_device.py +++ b/src/braket/pennylane_plugin/braket_device.py @@ -174,11 +174,7 @@ def operations(self) -> frozenset[str]: @property def observables(self) -> frozenset[str]: - base_observables = frozenset(super().observables) - # Amazon Braket only supports coefficients and multiple terms when shots==0 - if not self.shots: - return base_observables.union({"Hamiltonian", "LinearCombination"}) - return base_observables + return frozenset(super().observables) @property def circuit(self) -> Circuit: @@ -254,9 +250,8 @@ def _pl_to_braket_circuit( braket_circuit = self._apply_gradient_result_type(circuit, braket_circuit) elif not isinstance(circuit.measurements[0], MeasurementTransform): for measurement in circuit.measurements: - dev_wires = self.map_wires(measurement.wires).tolist() translated = translate_result_type( - measurement, dev_wires, self._braket_result_types + measurement.map_wires(self.wire_map), None, self._braket_result_types ) if isinstance(translated, tuple): for result_type in translated: @@ -281,7 +276,7 @@ def _apply_gradient_result_type(self, circuit, braket_circuit): f"Braket can only compute gradients for circuits with a single expectation" f" observable, not a {pl_measurements.return_type} observable." ) - if isinstance(pl_observable, (Hamiltonian, qml.Hamiltonian, Sum)): + if isinstance(pl_observable, (Hamiltonian, Sum)): targets = [self.map_wires(op.wires) for op in pl_observable.terms()[1]] else: targets = self.map_wires(pl_observable.wires).tolist() @@ -544,9 +539,10 @@ def _run_task(self, circuit, inputs=None): def _run_snapshots(self, snapshot_circuits, n_qubits, mapped_wires): raise NotImplementedError("Need to implement snapshots runner") - def _get_statistic(self, braket_result, observable): - dev_wires = self.map_wires(observable.wires).tolist() - return translate_result(braket_result, observable, dev_wires, self._braket_result_types) + 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]: @@ -641,6 +637,14 @@ def __init__( self._max_connections = max_connections self._max_retries = max_retries + @property + def observables(self) -> frozenset[str]: + base_observables = frozenset(super().observables) + # Amazon Braket only supports coefficients and multiple terms when shots==0 + if not self.shots: + return base_observables.union({"Hamiltonian", "LinearCombination"}) + return base_observables + @property def use_grouping(self) -> bool: # We *need* to do this because AdjointGradient doesn't support multiple diff --git a/src/braket/pennylane_plugin/translation.py b/src/braket/pennylane_plugin/translation.py index f75f5c0b..70716ea3 100644 --- a/src/braket/pennylane_plugin/translation.py +++ b/src/braket/pennylane_plugin/translation.py @@ -532,14 +532,16 @@ def get_adjoint_gradient_result_type( def translate_result_type( # noqa: C901 - measurement: MeasurementProcess, targets: list[int], supported_result_types: frozenset[str] + measurement: MeasurementProcess, + targets: Optional[list[int]], + supported_result_types: frozenset[str], ) -> Union[ResultType, tuple[ResultType, ...]]: """Translates a PennyLane ``MeasurementProcess`` into the corresponding Braket ``ResultType``. Args: measurement (MeasurementProcess): The PennyLane ``MeasurementProcess`` to translate - targets (list[int]): The target wires of the observable using a consecutive integer wire - ordering + targets (Optional[list[int]]): The target wires of the observable using a consecutive + integer wire ordering supported_result_types (frozenset[str]): Braket result types supported by the Braket device Returns: @@ -548,7 +550,9 @@ def translate_result_type( # noqa: C901 then this will return a result type for each term. """ return_type = measurement.return_type + targets = targets or measurement.wires.tolist() observable = measurement.obs + print(observable) if return_type is ObservableReturnTypes.Probability: return Probability(targets) @@ -560,25 +564,24 @@ def translate_result_type( # noqa: C901 return DensityMatrix(targets) raise NotImplementedError(f"Unsupported return type: {return_type}") - if isinstance(observable, (Hamiltonian, qml.Hamiltonian)): + if isinstance(observable, (Hamiltonian, qml.ops.Sum, qml.ops.SProd)): if return_type is ObservableReturnTypes.Expectation: - return tuple( - Expectation(_translate_observable(term), term.wires) for term in observable.ops - ) + simplified = qml.ops.LinearCombination(*observable.terms()).simplify() + return tuple(Expectation(_translate_observable(term)) for term in simplified.terms()[1]) raise NotImplementedError(f"Return type {return_type} unsupported for Hamiltonian") if observable is None: if return_type is ObservableReturnTypes.Counts: - return tuple(Sample(observables.Z(), target) for target in targets or measurement.wires) + return tuple(Sample(observables.Z(target)) for target in targets or measurement.wires) raise NotImplementedError(f"Unsupported return type: {return_type}") braket_observable = _translate_observable(observable) if return_type is ObservableReturnTypes.Expectation: - return Expectation(braket_observable, targets) + return Expectation(braket_observable) elif return_type is ObservableReturnTypes.Variance: - return Variance(braket_observable, targets) + return Variance(braket_observable) elif return_type in (ObservableReturnTypes.Sample, ObservableReturnTypes.Counts): - return Sample(braket_observable, targets) + return Sample(braket_observable) else: raise NotImplementedError(f"Unsupported return type: {return_type}") @@ -588,13 +591,12 @@ def _translate_observable(observable): raise qml.DeviceError(f"Unsupported observable: {type(observable)}") -@_translate_observable.register(Hamiltonian) -@_translate_observable.register(qml.Hamiltonian) -def _(H: Union[Hamiltonian, qml.Hamiltonian]): +@_translate_observable.register +def _(h: Hamiltonian): # terms is structured like [C, O] where C is a tuple of all the coefficients, and O is # a tuple of all the corresponding observable terms (X, Y, Z, H, etc or a tensor product # of them) - coefficents, pl_observables = H.terms() + coefficents, pl_observables = h.terms() braket_observables = list(map(lambda obs: _translate_observable(obs), pl_observables)) braket_hamiltonian = sum( (coef * obs for coef, obs in zip(coefficents[1:], braket_observables[1:])), @@ -604,33 +606,33 @@ def _(H: Union[Hamiltonian, qml.Hamiltonian]): @_translate_observable.register -def _(_: qml.PauliX): - return observables.X() +def _(obs: qml.PauliX): + return observables.X(obs.wires[0]) @_translate_observable.register -def _(_: qml.PauliY): - return observables.Y() +def _(obs: qml.PauliY): + return observables.Y(obs.wires[0]) @_translate_observable.register -def _(_: qml.PauliZ): - return observables.Z() +def _(obs: qml.PauliZ): + return observables.Z(obs.wires[0]) @_translate_observable.register -def _(_: qml.Hadamard): - return observables.H() +def _(obs: qml.Hadamard): + return observables.H(obs.wires[0]) @_translate_observable.register -def _(_: qml.Identity): - return observables.I() +def _(obs: qml.Identity): + return observables.I(obs.wires[0]) @_translate_observable.register -def _(h: qml.Hermitian): - return observables.Hermitian(qml.matrix(h)) +def _(obs: qml.Hermitian): + return observables.Hermitian(qml.matrix(obs), targets=obs.wires) _zero = np.array([[1, 0], [0, 0]]) @@ -638,15 +640,16 @@ def _(h: qml.Hermitian): @_translate_observable.register -def _(p: qml.Projector): - state, wires = p.parameters[0], p.wires +def _(obs: qml.Projector): + state = obs.parameters[0] + wires = obs.wires if len(state) == len(wires): # state is a basis state products = [_one if b else _zero for b in state] - hermitians = [observables.Hermitian(p) for p in products] + hermitians = [observables.Hermitian(p, targets=[w]) for p, w in zip(products, wires)] return observables.TensorProduct(hermitians) # state is a state vector - return observables.Hermitian(p.matrix()) + return observables.Hermitian(obs.matrix(), targets=wires) @_translate_observable.register @@ -672,7 +675,7 @@ def _(t: qml.ops.Sum): def translate_result( braket_result: GateModelQuantumTaskResult, measurement: MeasurementProcess, - targets: list[int], + targets: Optional[list[int]], supported_result_types: frozenset[str], ) -> Any: """Translates a Braket result into the corresponding PennyLane return type value. @@ -681,7 +684,7 @@ def translate_result( braket_result (GateModelQuantumTaskResult): The Braket result to translate. measurement (MeasurementProcess): The PennyLane measurement process associated with the result. - targets (list[int]): The qubits in the result. + targets (Optional[list[int]]): The qubits in the result. supported_result_types (frozenset[str]): The result types supported by the device. Returns: @@ -706,6 +709,7 @@ def translate_result( for i in sorted(key_indices) ] + targets = targets or measurement.wires.tolist() if measurement.return_type is ObservableReturnTypes.Counts and observable is None: if targets: new_dict = {} @@ -719,7 +723,7 @@ def translate_result( return dict(braket_result.measurement_counts) translated = translate_result_type(measurement, targets, supported_result_types) - if isinstance(observable, (Hamiltonian, qml.Hamiltonian)): + if isinstance(observable, (Hamiltonian, qml.ops.Sum, qml.ops.SProd)): coeffs, _ = observable.terms() return sum( coeff * braket_result.get_value_by_result_type(result_type) diff --git a/test/unit_tests/test_braket_device.py b/test/unit_tests/test_braket_device.py index 0cbbd07d..f62825a1 100644 --- a/test/unit_tests/test_braket_device.py +++ b/test/unit_tests/test_braket_device.py @@ -23,7 +23,7 @@ import pennylane as qml import pytest from braket.aws import AwsDevice, AwsDeviceType, AwsQuantumTask, AwsQuantumTaskBatch -from braket.circuits import Circuit, FreeParameter, Gate, Noise, Observable, result_types +from braket.circuits import Circuit, FreeParameter, Gate, Noise, observables, result_types from braket.circuits.noise_model import GateCriteria, NoiseModel, NoiseModelInstruction from braket.device_schema import DeviceActionType from braket.device_schema.gate_model_qpu_paradigm_properties_v1 import ( @@ -82,9 +82,9 @@ .i(2) .i(3) .probability(target=[0]) - .expectation(observable=Observable.X(), target=1) - .variance(observable=Observable.Y(), target=2) - .sample(observable=Observable.Z(), target=3) + .expectation(observable=observables.X(1)) + .variance(observable=observables.Y(2)) + .sample(observable=observables.Z(3)) ) DEVICE_ARN = "baz" @@ -196,17 +196,15 @@ def test_execute(mock_run): ) assert np.allclose( results[1], - RESULT.get_value_by_result_type( - result_types.Expectation(observable=Observable.X(), target=1) - ), + RESULT.get_value_by_result_type(result_types.Expectation(observable=observables.X(1))), ) assert np.allclose( results[2], - RESULT.get_value_by_result_type(result_types.Variance(observable=Observable.Y(), target=2)), + RESULT.get_value_by_result_type(result_types.Variance(observable=observables.Y(2))), ) assert np.allclose( results[3], - RESULT.get_value_by_result_type(result_types.Sample(observable=Observable.Z(), target=3)), + RESULT.get_value_by_result_type(result_types.Sample(observable=observables.Z(3))), ) assert dev.task == TASK EXPECTED_CIRC = ( @@ -218,9 +216,9 @@ def test_execute(mock_run): .i(2) .i(3) .probability(target=[0]) - .expectation(observable=Observable.X(), target=1) - .variance(observable=Observable.Y(), target=2) - .sample(observable=Observable.Z(), target=3) + .expectation(observable=observables.X(1)) + .variance(observable=observables.Y(2)) + .sample(observable=observables.Z(3)) ) mock_run.assert_called_with( EXPECTED_CIRC, @@ -255,17 +253,15 @@ def test_execute_parametrize_differentiable(mock_run): ) assert np.allclose( results[1], - RESULT.get_value_by_result_type( - result_types.Expectation(observable=Observable.X(), target=1) - ), + RESULT.get_value_by_result_type(result_types.Expectation(observable=observables.X(1))), ) assert np.allclose( results[2], - RESULT.get_value_by_result_type(result_types.Variance(observable=Observable.Y(), target=2)), + RESULT.get_value_by_result_type(result_types.Variance(observable=observables.Y(2))), ) assert np.allclose( results[3], - RESULT.get_value_by_result_type(result_types.Sample(observable=Observable.Z(), target=3)), + RESULT.get_value_by_result_type(result_types.Sample(observable=observables.Z(3))), ) assert dev.task == TASK EXPECTED_CIRC = ( @@ -279,9 +275,9 @@ def test_execute_parametrize_differentiable(mock_run): .i(2) .i(3) .probability(target=[0]) - .expectation(observable=Observable.X(), target=1) - .variance(observable=Observable.Y(), target=2) - .sample(observable=Observable.Z(), target=3) + .expectation(observable=observables.X(1)) + .variance(observable=observables.Y(2)) + .sample(observable=observables.Z(3)) ) mock_run.assert_called_with( EXPECTED_CIRC, @@ -378,7 +374,7 @@ def test_execute_parametrize_differentiable(mock_run): .cnot(0, 1) .rx(0, FreeParameter("p_0")) .ry(0, 0.543) - .adjoint_gradient(observable=Observable.X(), target=1, parameters=["p_0"]), + .adjoint_gradient(observable=observables.X(1), parameters=["p_0"]), 2, {"p_0": 0.432}, [ @@ -410,8 +406,7 @@ def test_execute_parametrize_differentiable(mock_run): .rx(0, FreeParameter("p_0")) .ry(0, FreeParameter("p_1")) .adjoint_gradient( - observable=(2 * Observable.X() @ Observable.Y()), - target=[0, 1], + observable=(2 * observables.X(0) @ observables.Y(1)), parameters=["p_0", "p_1"], ), 2, @@ -446,9 +441,9 @@ def test_execute_parametrize_differentiable(mock_run): .ry(0, FreeParameter("p_1")) .adjoint_gradient( observable=( - 2 * Observable.X() @ Observable.Y() + 0.75 * Observable.Y() @ Observable.Z() + 2 * observables.X(0) @ observables.Y(1) + + 0.75 * observables.Y(0) @ observables.Z(1) ), - target=[[0, 1], [0, 1]], parameters=["p_0", "p_1"], ), 2, @@ -521,8 +516,7 @@ def test_execute_with_gradient( .rx(0, FreeParameter("p_0")) .ry(0, FreeParameter("p_1")) .adjoint_gradient( - observable=(2 * Observable.X() @ Observable.Y()), - target=[0, 1], + observable=(2 * observables.X(0) @ observables.Y(1)), parameters=["p_0", "p_1"], ), 2, @@ -681,7 +675,7 @@ def test_pl_to_braket_circuit(): .rx(0, 0.2) .rx(1, 0.3) .cnot(0, 1) - .add_result_type(result_types.Expectation(observable=Observable.Z(), target=0)) + .add_result_type(result_types.Expectation(observable=observables.Z(0))) ) braket_circuit = dev._pl_to_braket_circuit(tape) @@ -706,9 +700,7 @@ def test_pl_to_braket_circuit_compute_gradient(): .rx(1, FreeParameter("p_1")) .cnot(0, 1) .add_result_type( - result_types.AdjointGradient( - observable=Observable.Z(), target=0, parameters=["p_0", "p_1"] - ) + result_types.AdjointGradient(observable=observables.Z(0), parameters=["p_0", "p_1"]) ) ) @@ -740,16 +732,14 @@ def test_pl_to_braket_circuit_compute_gradient_hamiltonian_tensor_product_terms( ) ) - braket_obs = 2 * Observable.X() @ Observable.X() + 3 * Observable.Y() @ Observable.Y() + braket_obs = 2 * observables.X(0) @ observables.X(1) + 3 * observables.Y(0) @ observables.Y(1) braket_circuit_true = ( Circuit() .rx(0, FreeParameter("p_0")) .rx(1, FreeParameter("p_1")) .cnot(0, 1) .add_result_type( - result_types.AdjointGradient( - observable=braket_obs, target=[[0, 1], [0, 1]], parameters=["p_0", "p_1"] - ) + result_types.AdjointGradient(observable=braket_obs, parameters=["p_0", "p_1"]) ) ) @@ -814,8 +804,8 @@ def test_pl_to_braket_circuit_hamiltonian(): .rx(0, 0.2) .rx(1, 0.3) .cnot(0, 1) - .expectation(Observable.X(), [0]) - .expectation(Observable.Y(), [1]) + .expectation(observables.X(0)) + .expectation(observables.Y(1)) ) braket_circuit = dev._pl_to_braket_circuit(tape) @@ -847,8 +837,8 @@ def test_pl_to_braket_circuit_hamiltonian_tensor_product_terms(): .rx(0, 0.2) .rx(1, 0.3) .cnot(0, 1) - .expectation(Observable.X() @ Observable.X(), [0, 1]) - .expectation(Observable.Y() @ Observable.Y(), [0, 1]) + .expectation(observables.X(0) @ observables.X(1)) + .expectation(observables.Y(0) @ observables.Y(1)) ) braket_circuit = dev._pl_to_braket_circuit(tape) @@ -935,21 +925,15 @@ def test_aws_device_batch_execute_parallel(mock_run_batch): ) assert np.allclose( results[1], - RESULT.get_value_by_result_type( - result_types.Expectation(observable=Observable.X(), target=1) - ), + RESULT.get_value_by_result_type(result_types.Expectation(observable=observables.X(1))), ) assert np.allclose( results[2], - RESULT.get_value_by_result_type( - result_types.Variance(observable=Observable.Y(), target=2) - ), + RESULT.get_value_by_result_type(result_types.Variance(observable=observables.Y(2))), ) assert np.allclose( results[3], - RESULT.get_value_by_result_type( - result_types.Sample(observable=Observable.Z(), target=3) - ), + RESULT.get_value_by_result_type(result_types.Sample(observable=observables.Z(3))), ) mock_run_batch.assert_called_with( @@ -991,21 +975,15 @@ def test_local_sim_batch_execute_parallel(mock_run_batch): ) assert np.allclose( results[1], - RESULT.get_value_by_result_type( - result_types.Expectation(observable=Observable.X(), target=1) - ), + RESULT.get_value_by_result_type(result_types.Expectation(observable=observables.X(1))), ) assert np.allclose( results[2], - RESULT.get_value_by_result_type( - result_types.Variance(observable=Observable.Y(), target=2) - ), + RESULT.get_value_by_result_type(result_types.Variance(observable=observables.Y(2))), ) assert np.allclose( results[3], - RESULT.get_value_by_result_type( - result_types.Sample(observable=Observable.Z(), target=3) - ), + RESULT.get_value_by_result_type(result_types.Sample(observable=observables.Z(3))), ) mock_run_batch.assert_called_with( @@ -1165,7 +1143,7 @@ def test_batch_execute_parametrize_differentiable(mock_run_batch): .cnot(0, 1) .i(2) .i(3) - .expectation(observable=Observable.X(), target=1) + .expectation(observable=observables.X(1)) ) expected_2 = ( @@ -1175,7 +1153,7 @@ def test_batch_execute_parametrize_differentiable(mock_run_batch): .cnot(0, 1) .i(2) .i(3) - .sample(observable=Observable.Z(), target=3) + .sample(observable=observables.Z(3)) ) circuits = [circuit1, circuit2] @@ -1477,6 +1455,7 @@ def test_unsupported_return_type(): mock_measurement.return_type = Enum("ObservableReturnTypes", {"Foo": "foo"}).Foo mock_measurement.obs = qml.PauliZ(0) mock_measurement.wires = qml.wires.Wires([0]) + mock_measurement.map_wires.return_value = mock_measurement tape = qml.tape.QuantumTape(measurements=[mock_measurement]) @@ -1696,7 +1675,7 @@ def test_add_braket_user_agent_invoked(aws_device_mock): .cnot(0, 1) .rx(0, FreeParameter("p_0")) .ry(0, 0.543) - .adjoint_gradient(observable=Observable.X(), target=1, parameters=["p_0"]), + .adjoint_gradient(observable=observables.X(1), parameters=["p_0"]), 2, {"p_0": 0.432}, [ @@ -1722,7 +1701,7 @@ def test_add_braket_user_agent_invoked(aws_device_mock): .cnot(0, 1) .rx(0, 0.432) .ry(0, 0.543) - .expectation(observable=Observable.X(), target=1), + .expectation(observable=observables.X(1)), 2, {}, [ @@ -1741,7 +1720,7 @@ def test_add_braket_user_agent_invoked(aws_device_mock): .rx(0, FreeParameter("p_1")) .unitary([0], 1 / np.sqrt(2) * anp.array([[1, 1], [1, -1]])) .cnot(0, 1) - .adjoint_gradient(observable=Observable.X(), target=1, parameters=["p_0"]), + .adjoint_gradient(observable=observables.X(1), parameters=["p_0"]), 2, {"p_1": 0.432}, [ @@ -1814,7 +1793,7 @@ def test_execute_and_gradients( .cnot(0, 1) .rx(0, 0.432) .ry(0, 0.543) - .variance(observable=Observable.X() @ Observable.Y(), target=[0, 1]), + .variance(observable=observables.X(0) @ observables.Y(1)), 2, {"p_1": 0.543}, [ @@ -2154,9 +2133,9 @@ def expected_braket_circuit_with_noise(): .i(2) .i(3) .probability(target=[0]) - .expectation(observable=Observable.X(), target=1) - .variance(observable=Observable.Y(), target=2) - .sample(observable=Observable.Z(), target=3) + .expectation(observable=observables.X(1)) + .variance(observable=observables.Y(2)) + .sample(observable=observables.Z(3)) ) @@ -2485,7 +2464,7 @@ def circuit(a): .add_verbatim_box( Circuit().gpi(0, x[0]).gpi2(0, x[0]).ms(0, 1, x[0], x[1]).ms(0, 1, x[0], x[1], x[2]) ) - .expectation(observable=Observable.Z(), target=1) + .expectation(observable=observables.Z(1)) ) mock_run.assert_called_with( expected_circuit, diff --git a/test/unit_tests/test_translation.py b/test/unit_tests/test_translation.py index d75aebc2..98e97f1b 100644 --- a/test/unit_tests/test_translation.py +++ b/test/unit_tests/test_translation.py @@ -831,22 +831,22 @@ def _result_meta() -> dict: "expected_braket_H, pl_H", [ ( - 2 * observables.X() @ observables.Y() @ observables.Z(), + 2 * observables.X(0) @ observables.Y(1) @ observables.Z(2), 2 * qml.PauliX(wires=0) @ qml.PauliY(wires=1) @ qml.PauliZ(wires=2), ), ( - 2 * (observables.X() @ observables.Y() @ observables.Z()), + 2 * (observables.X(0) @ observables.Y(1) @ observables.Z(2)), 2 * (qml.PauliX(wires=0) @ qml.PauliY(wires=1) @ qml.PauliZ(wires=2)), ), ( - 2 * observables.X() @ observables.Y() @ observables.Z() + 0.75 * observables.X(), + 2 * observables.X(0) @ observables.Y(1) @ observables.Z(2) + 0.75 * observables.X(0), 2 * qml.PauliX(wires=0) @ qml.PauliY(wires=1) @ qml.PauliZ(wires=2) + 0.75 * qml.PauliX(0), ), - (1.25 * observables.H(), 1.25 * qml.Hadamard(wires=0)), - (observables.X() @ observables.Y(), qml.ops.Prod(qml.PauliX(0), qml.PauliY(1))), - (observables.X() + observables.Y(), qml.ops.Sum(qml.PauliX(0), qml.PauliY(1))), - (observables.X(), qml.ops.SProd(scalar=4, base=qml.PauliX(0))), + (1.25 * observables.H(0), 1.25 * qml.Hadamard(wires=0)), + (observables.X(0) @ observables.Y(1), qml.ops.Prod(qml.PauliX(0), qml.PauliY(1))), + (observables.X(0) + observables.Y(1), qml.ops.Sum(qml.PauliX(0), qml.PauliY(1))), + (observables.X(0), qml.ops.SProd(scalar=4, base=qml.PauliX(0))), ], ) def test_translate_hamiltonian_observable(expected_braket_H, pl_H): From 039f92c5d35fbb808e099fca5d11670247be2ab7 Mon Sep 17 00:00:00 2001 From: Cody Wang Date: Thu, 5 Sep 2024 19:31:43 -0700 Subject: [PATCH 2/6] Update setup.py --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index aa4efcb5..b6db9d45 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.47.0", + "amazon-braket-sdk>=1.87.0", "autoray>=0.6.11", "pennylane>=0.34.0", ], From 42f8bea6df6f77c7b937e221e58cd5c5e98dff25 Mon Sep 17 00:00:00 2001 From: Cody Wang Date: Thu, 5 Sep 2024 21:29:28 -0700 Subject: [PATCH 3/6] Update translation.py --- src/braket/pennylane_plugin/translation.py | 43 ++++++++++------------ 1 file changed, 19 insertions(+), 24 deletions(-) diff --git a/src/braket/pennylane_plugin/translation.py b/src/braket/pennylane_plugin/translation.py index 70716ea3..2b1d84f3 100644 --- a/src/braket/pennylane_plugin/translation.py +++ b/src/braket/pennylane_plugin/translation.py @@ -34,7 +34,6 @@ from pennylane import numpy as np from pennylane.measurements import MeasurementProcess, ObservableReturnTypes from pennylane.operation import Observable, Operation -from pennylane.ops import Adjoint, Hamiltonian from pennylane.pulse import ParametrizedEvolution from braket.pennylane_plugin.ops import ( @@ -434,7 +433,7 @@ def _(ms: AAMS, parameters, device=None): @_translate_operation.register -def _(adjoint: Adjoint, parameters, device=None): +def _(adjoint: qml.ops.Adjoint, parameters, device=None): if isinstance(adjoint.base, qml.ISWAP): # gates.ISwap.adjoint() returns a different value return gates.PSwap(3 * np.pi / 2) @@ -523,8 +522,8 @@ def get_adjoint_gradient_result_type( ): if "AdjointGradient" not in supported_result_types: raise NotImplementedError("Unsupported return type: AdjointGradient") - braket_observable = _translate_observable(observable) + braket_observable = _translate_observable(_flatten_observable(observable)) braket_observable = ( braket_observable.item() if hasattr(braket_observable, "item") else braket_observable ) @@ -564,17 +563,18 @@ def translate_result_type( # noqa: C901 return DensityMatrix(targets) raise NotImplementedError(f"Unsupported return type: {return_type}") - if isinstance(observable, (Hamiltonian, qml.ops.Sum, qml.ops.SProd)): - if return_type is ObservableReturnTypes.Expectation: - simplified = qml.ops.LinearCombination(*observable.terms()).simplify() - return tuple(Expectation(_translate_observable(term)) for term in simplified.terms()[1]) - raise NotImplementedError(f"Return type {return_type} unsupported for Hamiltonian") - if observable is None: if return_type is ObservableReturnTypes.Counts: return tuple(Sample(observables.Z(target)) for target in targets or measurement.wires) raise NotImplementedError(f"Unsupported return type: {return_type}") + observable = _flatten_observable(observable) + + if isinstance(observable, qml.ops.LinearCombination): + if return_type is ObservableReturnTypes.Expectation: + return tuple(Expectation(_translate_observable(op)) for op in observable.terms()[1]) + raise NotImplementedError(f"Return type {return_type} unsupported for Hamiltonian") + braket_observable = _translate_observable(observable) if return_type is ObservableReturnTypes.Expectation: return Expectation(braket_observable) @@ -586,25 +586,19 @@ def translate_result_type( # noqa: C901 raise NotImplementedError(f"Unsupported return type: {return_type}") +def _flatten_observable(observable): + if isinstance(observable, (qml.ops.Hamiltonian, qml.ops.CompositeOp, qml.ops.SProd)): + simplified = qml.ops.LinearCombination(*observable.terms()).simplify() + coeffs, ops = simplified.terms() + return simplified if len(coeffs) > 1 or coeffs[0] != 1 else observable + return observable + + @singledispatch def _translate_observable(observable): raise qml.DeviceError(f"Unsupported observable: {type(observable)}") -@_translate_observable.register -def _(h: Hamiltonian): - # terms is structured like [C, O] where C is a tuple of all the coefficients, and O is - # a tuple of all the corresponding observable terms (X, Y, Z, H, etc or a tensor product - # of them) - coefficents, pl_observables = h.terms() - braket_observables = list(map(lambda obs: _translate_observable(obs), pl_observables)) - braket_hamiltonian = sum( - (coef * obs for coef, obs in zip(coefficents[1:], braket_observables[1:])), - coefficents[0] * braket_observables[0], - ) - return braket_hamiltonian - - @_translate_observable.register def _(obs: qml.PauliX): return observables.X(obs.wires[0]) @@ -723,7 +717,8 @@ def translate_result( return dict(braket_result.measurement_counts) translated = translate_result_type(measurement, targets, supported_result_types) - if isinstance(observable, (Hamiltonian, qml.ops.Sum, qml.ops.SProd)): + observable = _flatten_observable(observable) + if isinstance(observable, qml.ops.LinearCombination): coeffs, _ = observable.terms() return sum( coeff * braket_result.get_value_by_result_type(result_type) From 6a143a203edae52d8d5a908f5450d99318ac8754 Mon Sep 17 00:00:00 2001 From: Cody Wang Date: Fri, 6 Sep 2024 01:28:22 -0700 Subject: [PATCH 4/6] Update translation.py --- src/braket/pennylane_plugin/translation.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/braket/pennylane_plugin/translation.py b/src/braket/pennylane_plugin/translation.py index 2b1d84f3..17c8a226 100644 --- a/src/braket/pennylane_plugin/translation.py +++ b/src/braket/pennylane_plugin/translation.py @@ -589,8 +589,9 @@ def translate_result_type( # noqa: C901 def _flatten_observable(observable): if isinstance(observable, (qml.ops.Hamiltonian, qml.ops.CompositeOp, qml.ops.SProd)): simplified = qml.ops.LinearCombination(*observable.terms()).simplify() - coeffs, ops = simplified.terms() - return simplified if len(coeffs) > 1 or coeffs[0] != 1 else observable + coeffs, _ = simplified.terms() + if len(coeffs) > 1 or coeffs[0] != 1: + return simplified return observable From 3324dee24ef48a42b8643bc99f2d9b025b7aa870 Mon Sep 17 00:00:00 2001 From: Cody Wang Date: Fri, 6 Sep 2024 10:52:50 -0700 Subject: [PATCH 5/6] Update braket_device.py --- src/braket/pennylane_plugin/braket_device.py | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/src/braket/pennylane_plugin/braket_device.py b/src/braket/pennylane_plugin/braket_device.py index 48679b21..abfcfb4c 100644 --- a/src/braket/pennylane_plugin/braket_device.py +++ b/src/braket/pennylane_plugin/braket_device.py @@ -174,7 +174,11 @@ def operations(self) -> frozenset[str]: @property def observables(self) -> frozenset[str]: - return frozenset(super().observables) + base_observables = frozenset(super().observables) + # Amazon Braket only supports scalar multiplication and addition when shots==0 + if not self.shots: + return base_observables.union({"Hamiltonian", "LinearCombination"}) + return base_observables @property def circuit(self) -> Circuit: @@ -637,14 +641,6 @@ def __init__( self._max_connections = max_connections self._max_retries = max_retries - @property - def observables(self) -> frozenset[str]: - base_observables = frozenset(super().observables) - # Amazon Braket only supports coefficients and multiple terms when shots==0 - if not self.shots: - return base_observables.union({"Hamiltonian", "LinearCombination"}) - return base_observables - @property def use_grouping(self) -> bool: # We *need* to do this because AdjointGradient doesn't support multiple From 15012c2462beee83b459d98daa592c4c5ac3ce6c Mon Sep 17 00:00:00 2001 From: Cody Wang Date: Fri, 6 Sep 2024 13:02:17 -0700 Subject: [PATCH 6/6] Update src/braket/pennylane_plugin/translation.py Co-authored-by: Abe Coull <85974725+math411@users.noreply.github.com> --- src/braket/pennylane_plugin/translation.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/braket/pennylane_plugin/translation.py b/src/braket/pennylane_plugin/translation.py index 17c8a226..7d1fc2ac 100644 --- a/src/braket/pennylane_plugin/translation.py +++ b/src/braket/pennylane_plugin/translation.py @@ -551,7 +551,6 @@ def translate_result_type( # noqa: C901 return_type = measurement.return_type targets = targets or measurement.wires.tolist() observable = measurement.obs - print(observable) if return_type is ObservableReturnTypes.Probability: return Probability(targets)