diff --git a/setup.py b/setup.py index 6947c19a..e0218975 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.35.0", + "amazon-braket-sdk>=1.47.0", "pennylane==0.30.0", ], entry_points={ diff --git a/src/braket/pennylane_plugin/__init__.py b/src/braket/pennylane_plugin/__init__.py index c4f45b1c..1e9c01ff 100644 --- a/src/braket/pennylane_plugin/__init__.py +++ b/src/braket/pennylane_plugin/__init__.py @@ -20,6 +20,7 @@ BraketLocalQubitDevice, ) from braket.pennylane_plugin.ops import ( # noqa: F401 + AAMS, MS, PSWAP, CPhaseShift00, diff --git a/src/braket/pennylane_plugin/ops.py b/src/braket/pennylane_plugin/ops.py index 83c0c2cd..3b09c90a 100644 --- a/src/braket/pennylane_plugin/ops.py +++ b/src/braket/pennylane_plugin/ops.py @@ -410,7 +410,7 @@ class MS(Operation): IonQ native Mølmer-Sørenson gate. - .. math:: \mathtt{MS}(\phi_0, \phi_1) = \begin{bmatrix} + .. math:: \mathtt{MS}(\phi_0, \phi_1) = \frac{1}{\sqrt{2}}\begin{bmatrix} 1 & 0 & 0 & -ie^{-i (\phi_0 + \phi_1)} \\ 0 & 1 & -ie^{-i (\phi_0 - \phi_1)} & 0 \\ 0 & -ie^{i (\phi_0 - \phi_1)} & 1 & 0 \\ @@ -456,3 +456,60 @@ def compute_matrix(phi_0, phi_1): def adjoint(self): (phi_0, phi_1) = self.parameters return MS(phi_0 + np.pi, phi_1, wires=self.wires) + + +class AAMS(Operation): + r""" AAMS(phi_0, phi_1, theta, wires) + + IonQ native Arbitrary-Angle Mølmer-Sørenson gate. + + + .. math:: \mathtt{MS}(\phi_0, \phi_1, \theta) = \begin{bmatrix} + \cos{\frac{\theta}{2}} & 0 & 0 & -ie^{-i (\phi_0 + \phi_1)}\sin{\frac{\theta}{2}} \\ + 0 & \cos{\frac{\theta}{2}} & -ie^{-i (\phi_0 - \phi_1)}\sin{\frac{\theta}{2}} & 0 \\ + 0 & -ie^{i (\phi_0 - \phi_1)}\sin{\frac{\theta}{2}} & \cos{\frac{\theta}{2}} & 0 \\ + -ie^{i (\phi_0 + \phi_1)}\sin{\frac{\theta}{2}} & 0 & 0 & \cos{\frac{\theta}{2}} + \end{bmatrix}. + + **Details:** + + * Number of wires: 2 + * Number of parameters: 2 + + Args: + phi_0 (float): the first phase angle + phi_1 (float): the second phase angle + theta (float): the entangling angle + wires (int): the subsystem the gate acts on + do_queue (bool): Indicates whether the operator should be + immediately pushed into the Operator queue (optional) + id (str or None): String representing the operation (optional) + """ + num_params = 3 + num_wires = 2 + grad_method = "F" + + def __init__(self, phi_0, phi_1, theta, wires, do_queue=True, id=None): + super().__init__(phi_0, phi_1, theta, wires=wires, do_queue=do_queue, id=id) + + @staticmethod + def compute_matrix(phi_0, phi_1, theta): + if qml.math.get_interface(phi_0) == "tensorflow": + phi_0 = qml.math.cast_like(phi_0, 1j) + if qml.math.get_interface(phi_1) == "tensorflow": + phi_1 = qml.math.cast_like(phi_1, 1j) + if qml.math.get_interface(theta) == "tensorflow": + theta = qml.math.cast_like(theta, 1j) + + return np.array( + [ + [np.cos(theta / 2), 0, 0, -1j * np.exp(-1j * (phi_0 + phi_1)) * np.sin(theta / 2)], + [0, np.cos(theta / 2), -1j * np.exp(-1j * (phi_0 - phi_1)) * np.sin(theta / 2), 0], + [0, -1j * np.exp(1j * (phi_0 - phi_1)) * np.sin(theta / 2), np.cos(theta / 2), 0], + [-1j * np.exp(1j * (phi_0 + phi_1)) * np.sin(theta / 2), 0, 0, np.cos(theta / 2)], + ] + ) + + def adjoint(self): + (phi_0, phi_1, theta) = self.parameters + return AAMS(phi_0 + np.pi, phi_1, theta, wires=self.wires) diff --git a/src/braket/pennylane_plugin/translation.py b/src/braket/pennylane_plugin/translation.py index 94e8d5e4..1b2d6847 100644 --- a/src/braket/pennylane_plugin/translation.py +++ b/src/braket/pennylane_plugin/translation.py @@ -33,6 +33,7 @@ from pennylane.ops import Adjoint from braket.pennylane_plugin.ops import ( + AAMS, MS, PSWAP, CPhaseShift00, @@ -85,7 +86,7 @@ "ecr": "ECR", "gpi": "GPi", "gpi2": "GPi2", - "ms": "MS", + "ms": "AAMS", } @@ -104,11 +105,15 @@ def supported_operations(device: Device) -> FrozenSet[str]: raise AttributeError("Device needs to have properties defined.") supported_ops = frozenset(op.lower() for op in properties.supportedOperations) supported_pragmas = frozenset(op.lower() for op in properties.supportedPragmas) - return frozenset( + translated = frozenset( _BRAKET_TO_PENNYLANE_OPERATIONS[op] for op in _BRAKET_TO_PENNYLANE_OPERATIONS if op.lower() in supported_ops or f"braket_noise_{op.lower()}" in supported_pragmas ) + # both AAMS and MS map to ms + if "AAMS" in translated: + translated |= {"MS"} + return translated def translate_operation( @@ -390,6 +395,12 @@ def _(ms: MS, parameters): return gates.MS(phi_0, phi_1) +@_translate_operation.register +def _(ms: AAMS, parameters): + phi_0, phi_1, theta = parameters[:3] + return gates.MS(phi_0, phi_1, theta) + + @_translate_operation.register def _(adjoint: Adjoint, parameters): if isinstance(adjoint.base, qml.ISWAP): diff --git a/test/unit_tests/test_ops.py b/test/unit_tests/test_ops.py index 38c998a5..a405eb68 100644 --- a/test/unit_tests/test_ops.py +++ b/test/unit_tests/test_ops.py @@ -26,7 +26,7 @@ from numpy import float64 from braket.pennylane_plugin import PSWAP, CPhaseShift00, CPhaseShift01, CPhaseShift10 -from braket.pennylane_plugin.ops import MS, GPi, GPi2 +from braket.pennylane_plugin.ops import AAMS, MS, GPi, GPi2 gates_1q_parametrized = [ (GPi, gates.GPi), @@ -44,6 +44,10 @@ (MS, gates.MS), ] +gates_2q_3p_parametrized = [ + (AAMS, gates.MS), +] + gates_2q_non_parametrized = [] # Empty... For now! observables_1q = [ @@ -62,7 +66,10 @@ def test_ops_parametrized(pl_op, braket_gate, angle): _assert_decomposition(pl_op, params=[angle]) -@pytest.mark.parametrize("pl_op, braket_gate", gates_2q_parametrized) +@pytest.mark.parametrize( + "pl_op, braket_gate", + gates_2q_parametrized, +) @pytest.mark.parametrize( "angle", [tf.Variable(((i + 1) * math.pi / 12), dtype=float64) for i in range(12)] ) @@ -70,27 +77,43 @@ def test_ops_parametrized_tf(pl_op, braket_gate, angle): """Tests that the matrices and decompositions of parametrized custom operations are correct using tensorflow interface. """ - pl_op.compute_matrix(angle) - _assert_decomposition(pl_op, params=[angle]) + angles = [angle] * pl_op.num_params + pl_op.compute_matrix(*angles) + _assert_decomposition(pl_op, params=angles) -@pytest.mark.parametrize("pl_op, braket_gate", gates_1q_parametrized + gates_2q_2p_parametrized) -@pytest.mark.parametrize("angle", [(i + 1) * math.pi / 12 for i in range(12)]) -def test_ops_parametrized_no_decomposition(pl_op, braket_gate, angle): +@pytest.mark.parametrize( + "pl_op, braket_gate", + gates_1q_parametrized + gates_2q_2p_parametrized + gates_2q_3p_parametrized, +) +@pytest.mark.parametrize("angle_1", [(i + 1) * math.pi / 12 for i in range(12)]) +@pytest.mark.parametrize("angle_2", [(i + 1) * math.pi / 12 for i in range(12)]) +@pytest.mark.parametrize("angle_3", [(i + 1) * math.pi / 12 for i in range(6)]) +def test_ops_parametrized_no_decomposition(pl_op, braket_gate, angle_1, angle_2, angle_3): """Tests that the matrices and decompositions of parametrized custom operations are correct.""" - angles = [angle] * pl_op.num_params + angles = [angle_1, angle_2, angle_3][: pl_op.num_params] assert np.allclose(pl_op.compute_matrix(*angles), braket_gate(*angles).to_matrix()) -@pytest.mark.parametrize("pl_op, braket_gate", gates_1q_parametrized + gates_2q_2p_parametrized) @pytest.mark.parametrize( - "angle", [tf.Variable(((i + 1) * math.pi / 12), dtype=float64) for i in range(12)] + "pl_op, braket_gate", + gates_1q_parametrized + gates_2q_2p_parametrized + gates_2q_3p_parametrized, ) -def test_ops_parametrized_tf_no_decomposition(pl_op, braket_gate, angle): +@pytest.mark.parametrize( + "angle_1", [tf.Variable(((i + 1) * math.pi / 12), dtype=float64) for i in range(12)] +) +@pytest.mark.parametrize( + "angle_2", [tf.Variable(((i + 1) * math.pi / 12), dtype=float64) for i in range(12)] +) +@pytest.mark.parametrize( + "angle_3", [tf.Variable(((i + 1) * math.pi / 12), dtype=float64) for i in range(6)] +) +def test_ops_parametrized_tf_no_decomposition(pl_op, braket_gate, angle_1, angle_2, angle_3): """Tests that the matrices and decompositions of parametrized custom operations are correct using tensorflow interface. """ - pl_op.compute_matrix(*[angle] * pl_op.num_params) + angles = [angle_1, angle_2, angle_3][: pl_op.num_params] + pl_op.compute_matrix(*angles) @pytest.mark.parametrize("pl_op, braket_gate", gates_2q_non_parametrized) @@ -103,7 +126,10 @@ def test_ops_non_parametrized(pl_op, braket_gate): @patch("braket.pennylane_plugin.ops.np", new=anp) -@pytest.mark.parametrize("pl_op, braket_gate", gates_2q_parametrized) +@pytest.mark.parametrize( + "pl_op, braket_gate", + gates_2q_parametrized, +) @pytest.mark.parametrize("angle", [(i + 1) * math.pi / 12 for i in range(12)]) @pytest.mark.parametrize("observable", observables_2q) def test_param_shift_2q(pl_op, braket_gate, angle, observable): @@ -183,7 +209,11 @@ def _assert_decomposition(pl_op, params=None): @pytest.mark.parametrize( - "pl_op, braket_gate", gates_1q_parametrized + gates_2q_parametrized + gates_2q_2p_parametrized + "pl_op, braket_gate", + gates_1q_parametrized + + gates_2q_parametrized + + gates_2q_2p_parametrized + + gates_2q_3p_parametrized, ) @pytest.mark.parametrize("angle", [(i + 1) * math.pi / 12 for i in range(12)]) def test_gate_adjoint_parametrized(pl_op, braket_gate, angle): @@ -204,7 +234,10 @@ def test_gate_adjoint_non_parametrized(pl_op, braket_gate): ) -@pytest.mark.parametrize("pl_op, braket_gate", gates_2q_parametrized) +@pytest.mark.parametrize( + "pl_op, braket_gate", + gates_2q_parametrized, +) @pytest.mark.parametrize("angle", [(i + 1) * math.pi / 12 for i in range(12)]) def test_gate_generator(pl_op, braket_gate, angle): op = pl_op(angle, wires=[0, 1]) diff --git a/test/unit_tests/test_shadow_expval.py b/test/unit_tests/test_shadow_expval.py index 2fd544a6..b0331286 100644 --- a/test/unit_tests/test_shadow_expval.py +++ b/test/unit_tests/test_shadow_expval.py @@ -1,22 +1,22 @@ import json +from typing import Any, Dict, Optional from unittest import mock from unittest.mock import Mock, PropertyMock, patch -from typing import Any, Dict, Optional +import braket.ir as ir import pennylane as qml import pytest -import braket.ir as ir 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 OpenQASMDeviceActionProperties -from braket.simulator import BraketSimulator from braket.device_schema.simulators import GateModelSimulatorDeviceCapabilities from braket.devices import LocalSimulator +from braket.simulator import BraketSimulator from braket.task_result import GateModelTaskResult from braket.tasks import GateModelQuantumTaskResult -from pennylane.tape import QuantumScript, QuantumTape from pennylane.measurements import MeasurementTransform +from pennylane.tape import QuantumScript, QuantumTape from pennylane.wires import Wires from braket.pennylane_plugin import BraketAwsQubitDevice, BraketLocalQubitDevice diff --git a/test/unit_tests/test_translation.py b/test/unit_tests/test_translation.py index b44ff268..fa6245e5 100644 --- a/test/unit_tests/test_translation.py +++ b/test/unit_tests/test_translation.py @@ -35,7 +35,7 @@ from pennylane.wires import Wires from braket.pennylane_plugin import PSWAP, CPhaseShift00, CPhaseShift01, CPhaseShift10 -from braket.pennylane_plugin.ops import MS, GPi, GPi2 +from braket.pennylane_plugin.ops import AAMS, MS, GPi, GPi2 from braket.pennylane_plugin.translation import ( _BRAKET_TO_PENNYLANE_OPERATIONS, _translate_observable, @@ -72,6 +72,7 @@ (GPi, gates.GPi, [0], [2]), (GPi2, gates.GPi2, [0], [2]), (MS, gates.MS, [0, 1], [2, 3]), + (AAMS, gates.MS, [0, 1], [2, 3, 0.5]), (qml.ECR, gates.ECR, [0, 1], []), (qml.ISWAP, gates.ISwap, [0, 1], []), (PSWAP, gates.PSwap, [0, 1], [np.pi]), @@ -137,6 +138,7 @@ (GPi, gates.GPi, [0], [2], [2]), (GPi2, gates.GPi2, [0], [2], [2 + np.pi]), (MS, gates.MS, [0, 1], [2, 3], [2 + np.pi, 3]), + (AAMS, gates.MS, [0, 1], [2, 3, 0.5], [2 + np.pi, 3, 0.5]), (PSWAP, gates.PSwap, [0, 1], [0.15], [-0.15]), (qml.IsingXX, gates.XX, [0, 1], [0.15], [-0.15]), (qml.IsingXY, gates.XY, [0, 1], [0.15], [-0.15]), @@ -178,6 +180,14 @@ (GPi, gates.GPi, [0], [2], ["a"], [FreeParameter("a")]), (GPi2, gates.GPi2, [0], [2], ["a"], [FreeParameter("a")]), (MS, gates.MS, [0, 1], [2, 3], ["a", "b"], [FreeParameter("a"), FreeParameter("b")]), + ( + AAMS, + gates.MS, + [0, 1], + [2, 3, 0.5], + ["a", "b", "c"], + [FreeParameter("a"), FreeParameter("b"), FreeParameter("c")], + ), (PSWAP, gates.PSwap, [0, 1], [np.pi], ["pi"], [FreeParameter("pi")]), (qml.ECR, gates.ECR, [0, 1], [], [], []), (qml.ISWAP, gates.ISwap, [0, 1], [], [], []), @@ -255,12 +265,16 @@ def test_translate_operation(pl_cls, braket_cls, qubits, params): pl_op = pl_cls(*params, wires=qubits) braket_gate = braket_cls(*params) assert translate_operation(pl_op) == braket_gate - if isinstance(pl_op, (GPi, GPi2, MS)): + if isinstance(pl_op, (GPi, GPi2, MS, AAMS)): + translated_back = _braket_to_pl[ + re.match("^[a-z0-2]+", braket_gate.to_ir(qubits, ir_type=IRType.OPENQASM)).group(0) + ] assert ( - _braket_to_pl[ - re.match("^[a-z0-2]+", braket_gate.to_ir(qubits, ir_type=IRType.OPENQASM)).group(0) - ] - == pl_op.name + translated_back == pl_op.name + if pl_op.name != "MS" + # PL MS and AAMS both get translated to Braket MS. + # Braket MS gets translated to PL AAMS. + else translated_back == "AAMS" ) else: assert ( @@ -283,12 +297,16 @@ def test_translate_operation_with_unique_params( translate_operation(pl_op, use_unique_params=True, param_names=pl_param_names) == braket_gate ) - if isinstance(pl_op, (GPi, GPi2, MS)): + if isinstance(pl_op, (GPi, GPi2, MS, AAMS)): + translated_back = _braket_to_pl[ + re.match("^[a-z0-2]+", braket_gate.to_ir(qubits, ir_type=IRType.OPENQASM)).group(0) + ] assert ( - _braket_to_pl[ - re.match("^[a-z0-2]+", braket_gate.to_ir(qubits, ir_type=IRType.OPENQASM)).group(0) - ] - == pl_op.name + translated_back == pl_op.name + if pl_op.name != "MS" + # PL MS and AAMS both get translated to Braket MS. + # Braket MS gets translated to PL AAMS. + else translated_back == "AAMS" ) else: assert ( @@ -303,7 +321,7 @@ def test_translate_operation_inverse(pl_cls, braket_cls, qubits, params, inv_par pl_op = qml.adjoint(pl_cls(*params, wires=qubits)) braket_gate = braket_cls(*inv_params) assert translate_operation(pl_op) == braket_gate - if isinstance(pl_op.base, (GPi, GPi2, MS)): + if isinstance(pl_op.base, (GPi, GPi2, MS, AAMS)): op_name = _braket_to_pl[ re.match( "^[a-z0-2]+", @@ -315,7 +333,14 @@ def test_translate_operation_inverse(pl_cls, braket_cls, qubits, params, inv_par braket_gate.to_ir(qubits).__class__.__name__.lower().replace("_", "") ] - assert f"Adjoint({op_name})" == pl_op.name + assert ( + f"Adjoint({op_name})" == pl_op.name + if pl_op.name != "Adjoint(MS)" + # PL MS and AAMS both get translated to Braket MS. + # Braket MS gets translated to PL AAMS. + else f"Adjoint({op_name})" == "Adjoint(AAMS)" + ) + # assert f"Adjoint({op_name})" == pl_op.name @patch("braket.circuits.gates.X.adjoint")