feat: arbitrary angle MS gate (#174)

Co-authored-by: Cody Wang <[email protected]>

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={

src/braket/pennylane_plugin/__init__.py

@@ -20,6 +20,7 @@
     BraketLocalQubitDevice,
 )
 from braket.pennylane_plugin.ops import (  # noqa: F401
+    AAMS,
     MS,
     PSWAP,
     CPhaseShift00,

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)

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):

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,35 +66,54 @@ 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)]
 )
 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])

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

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")