fix: Use PL GeneralizedAmplitudeDamping convention (#322)

Author: speller26

setup.py

@@ -29,7 +29,7 @@
     install_requires=[
         "amazon-braket-sdk>=1.97.0",
         "autoray>=0.6.11",
-        "pennylane>=0.42.0",
+        "pennylane>=0.44.0",
     ],
     entry_points={
         "pennylane.plugins": [

src/braket/pennylane_plugin/translation.py

@@ -319,7 +319,7 @@ def _(_: qml.AmplitudeDamping, parameters, device=None):
 @_translate_operation.register
 def _(_: qml.GeneralizedAmplitudeDamping, parameters, device=None):
     gamma = parameters[0]
-    probability = parameters[1]
+    probability = 1 - parameters[1]
     return noises.GeneralizedAmplitudeDamping(probability=probability, gamma=gamma)
 
 

test/integ_tests/test_apply.py

@@ -154,8 +154,8 @@ def test_single_qubit_noise(self, init_state, dm_device, op, prob, tol):
         state = init_state(1)
         TestHardwareApply.assert_noise_op(op, dev, state, [0], tol, [prob])
 
-    @pytest.mark.parametrize("gamma", [0.0, 0.42])
     @pytest.mark.parametrize("prob", [0.0, 0.42])
+    @pytest.mark.parametrize("gamma", [0.0, 0.42])
     def test_generalized_amplitude_damping(self, init_state, dm_device, gamma, prob, tol):
         """Test parametrized GeneralizedAmplitudeDamping"""
         dev = dm_device(1)

test/unit_tests/test_braket_device.py

@@ -239,17 +239,17 @@ def test_apply_unique_parameters():
             qml.RX(np.pi, wires=0),
             qml.RY(np.pi, wires=0),
             # note the gamma/p ordering doesn't affect the naming of the parameters below.
-            qml.GeneralizedAmplitudeDamping(gamma=0.1, p=0.9, wires=0),
-            qml.GeneralizedAmplitudeDamping(p=0.9, gamma=0.1, wires=0),
+            qml.GeneralizedAmplitudeDamping(gamma=0.1, p=0.2, wires=0),
+            qml.GeneralizedAmplitudeDamping(p=0.2, gamma=0.1, wires=0),
         ],
         use_unique_params=True,
     )
     expected = Circuit().h(0).cnot(0, 1).rx(0, FreeParameter("p_0"))
     expected = expected.ry(0, FreeParameter("p_1"))
 
     # Right now, the Braket SDK doesn't keep track of noise parameters
-    expected = expected.generalized_amplitude_damping(0, gamma=0.1, probability=0.9)
-    expected = expected.generalized_amplitude_damping(0, gamma=0.1, probability=0.9)
+    expected = expected.generalized_amplitude_damping(0, gamma=0.1, probability=0.8)
+    expected = expected.generalized_amplitude_damping(0, gamma=0.1, probability=0.8)
     assert circuit == expected
 
 

test/unit_tests/test_translation.py

@@ -20,7 +20,7 @@
 import pennylane as qml
 import pytest
 from braket.aws import AwsDevice, AwsDeviceType
-from braket.circuits import FreeParameter, gates, noises, observables
+from braket.circuits import FreeParameter, Noise, gates, noises, observables
 from braket.circuits.result_types import (
     AdjointGradient,
     DensityMatrix,
@@ -151,12 +151,6 @@ def _aws_device(
     (qml.IsingYY, gates.YY, [0, 1], [np.pi]),
     (qml.IsingZZ, gates.ZZ, [0, 1], [np.pi]),
     (qml.AmplitudeDamping, noises.AmplitudeDamping, [0], [0.1]),
-    (
-        qml.GeneralizedAmplitudeDamping,
-        noises.GeneralizedAmplitudeDamping,
-        [0],
-        [0.1, 0.15],
-    ),
     (qml.PhaseDamping, noises.PhaseDamping, [0], [0.1]),
     (qml.DepolarizingChannel, noises.Depolarizing, [0], [0.1]),
     (qml.BitFlip, noises.BitFlip, [0], [0.1]),
@@ -315,14 +309,6 @@ def _aws_device(
         ["alpha"],
         [FreeParameter("alpha")],
     ),
-    (
-        qml.GeneralizedAmplitudeDamping,
-        noises.GeneralizedAmplitudeDamping,
-        [0],
-        [0.1, 0.15],
-        ["p_000", "p_001"],
-        [FreeParameter("p_000"), FreeParameter("p_001")],
-    ),
     (qml.PhaseDamping, noises.PhaseDamping, [0], [0.1], ["a"], [FreeParameter("a")]),
     (
         qml.DepolarizingChannel,
@@ -386,7 +372,12 @@ 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, AAMS, PRx)):
+    if isinstance(braket_gate, (Noise, gates.Unitary)):
+        assert (
+            _braket_to_pl[braket_gate.to_ir(qubits).__class__.__name__.lower().replace("_", "")]
+            == pl_op.name
+        )
+    else:
         translated_back = _braket_to_pl[
             re.match("^[a-z0-2]+", braket_gate.to_ir(qubits, ir_type=IRType.OPENQASM)).group(0)
         ]
@@ -397,11 +388,6 @@ def test_translate_operation(pl_cls, braket_cls, qubits, params):
             # Braket MS gets translated to PL AAMS.
             else translated_back == "AAMS"
         )
-    else:
-        assert (
-            _braket_to_pl[braket_gate.to_ir(qubits).__class__.__name__.lower().replace("_", "")]
-            == pl_op.name
-        )
 
 
 @pytest.mark.parametrize(
@@ -418,7 +404,12 @@ 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, AAMS)):
+    if isinstance(braket_gate, (Noise, gates.Unitary)):
+        assert (
+            _braket_to_pl[braket_gate.to_ir(qubits).__class__.__name__.lower().replace("_", "")]
+            == pl_op.name
+        )
+    else:
         translated_back = _braket_to_pl[
             re.match("^[a-z0-2]+", braket_gate.to_ir(qubits, ir_type=IRType.OPENQASM)).group(0)
         ]
@@ -429,11 +420,32 @@ def test_translate_operation_with_unique_params(
             # Braket MS gets translated to PL AAMS.
             else translated_back == "AAMS"
         )
-    else:
-        assert (
-            _braket_to_pl[braket_gate.to_ir(qubits).__class__.__name__.lower().replace("_", "")]
-            == pl_op.name
-        )
+
+
+def test_generalized_amplitude_damping():
+    """Tests that GeneralizedAmplitudeDamping is translated correctly"""
+    qubits = [0]
+    pl_param_names = ["p_000", "p_001"]
+    pl_op = qml.GeneralizedAmplitudeDamping(0.1, 0.15, wires=qubits)
+    braket_op = noises.GeneralizedAmplitudeDamping(0.1, 0.85)
+    assert translate_operation(pl_op) == braket_op
+    assert (
+        _braket_to_pl[braket_op.to_ir(qubits).__class__.__name__.lower().replace("_", "")]
+        == pl_op.name
+    )
+    braket_op_parametrized = noises.GeneralizedAmplitudeDamping(
+        FreeParameter("p_000"), 1 - FreeParameter("p_001")
+    )
+    assert (
+        translate_operation(pl_op, use_unique_params=True, param_names=pl_param_names)
+        == braket_op_parametrized
+    )
+    assert (
+        _braket_to_pl[
+            braket_op_parametrized.to_ir(qubits).__class__.__name__.lower().replace("_", "")
+        ]
+        == pl_op.name
+    )
 
 
 def amplitude(p, t):