Add U gate

src/braket/circuits/gates.py

@@ -1091,6 +1091,118 @@ def phaseshift(
 Gate.register_gate(PhaseShift)
 
 
+class U(TripleAngledGate):
+    """Parameterized primitive gate for OpenQASM simulator
+
+    Args:
+        angle_1 (Union[FreeParameterExpression, float]): theta angle in radians.
+        angle_2 (Union[FreeParameterExpression, float]): phi angle in radians.
+        angle_3 (Union[FreeParameterExpression, float]): lambda angle in radians.
+    """
+
+    def __init__(
+        self,
+        angle_1: Union[FreeParameterExpression, float],
+        angle_2: Union[FreeParameterExpression, float],
+        angle_3: Union[FreeParameterExpression, float],
+    ):
+        super().__init__(
+            angle_1=angle_1,
+            angle_2=angle_2,
+            angle_3=angle_3,
+            qubit_count=None,
+            ascii_symbols=[_multi_angled_ascii_characters("U", angle_1, angle_2, angle_3)],
+        )
+
+    @property
+    def _qasm_name(self) -> str:
+        return "U"
+
+    def to_matrix(self) -> np.ndarray:
+        """
+        Generate parameterized Unitary matrix.
+        https://openqasm.com/language/gates.html#built-in-gates
+
+        Returns:
+            np.ndarray: U Matrix
+        """
+        theta = self.angle_1
+        phi = self.angle_2
+        lam = self.angle_3
+        return np.array(
+            [
+                [
+                    np.cos(theta / 2),
+                    -np.exp(1j * lam) * np.sin(theta / 2),
+                ],
+                [
+                    np.exp(1j * phi) * np.sin(theta / 2),
+                    np.exp(1j * (phi + lam)) * np.cos(theta / 2),
+                ],
+            ]
+        )
+
+    def adjoint(self) -> list[Gate]:
+        return [U(self.angle_1, np.pi - self.angle_3, np.pi - self.angle_2)]
+
+    @staticmethod
+    def fixed_qubit_count() -> int:
+        return 1
+
+    def bind_values(self, **kwargs) -> U:
+        return _get_angles(self, **kwargs)
+
+    @staticmethod
+    @circuit.subroutine(register=True)
+    def u(
+        target: QubitInput,
+        angle_1: Union[FreeParameterExpression, float],
+        angle_2: Union[FreeParameterExpression, float],
+        angle_3: Union[FreeParameterExpression, float],
+        *,
+        control: Optional[QubitSetInput] = None,
+        control_state: Optional[BasisStateInput] = None,
+        power: float = 1,
+    ) -> Iterable[Instruction]:
+        """Registers this function into the circuit class.
+
+        Args:
+            target1 (QubitInput): Target qubit 1 index.
+            angle_1 (Union[FreeParameterExpression, float]): theta angle in radians.
+            angle_2 (Union[FreeParameterExpression, float]): phi angle in radians.
+            angle_3 (Union[FreeParameterExpression, float]): lambda angle in radians.
+            control (Optional[QubitSetInput]): Control qubit(s). Default None.
+            control_state (Optional[BasisStateInput]): Quantum state on which to control the
+                operation. Must be a binary sequence of same length as number of qubits in
+                `control`. Will be ignored if `control` is not present. May be represented as a
+                string, list, or int. For example "0101", [0, 1, 0, 1], 5 all represent
+                controlling on qubits 0 and 2 being in the \\|0⟩ state and qubits 1 and 3 being
+                in the \\|1⟩ state. Default "1" * len(control).
+            power (float): Integer or fractional power to raise the gate to. Negative
+                powers will be split into an inverse, accompanied by the positive power.
+                Default 1.
+
+        Returns:
+            Iterable[Instruction]: MS instruction.
+
+        Examples:
+            >>> circ = Circuit().ms(0, 1, 0.15, 0.34)
+        """
+        return [
+            Instruction(
+                U(angle_1, angle_2, angle_3),
+                target=qubit,
+                control=control,
+                control_state=control_state,
+                power=power,
+            )
+            for qubit in QubitSet(target)
+        ]
+
+
+Gate.register_gate(U)
+
+
 # Two qubit gates #
 
 

src/braket/circuits/translations.py

@@ -55,6 +55,7 @@
     "rx": braket_gates.Rx,
     "ry": braket_gates.Ry,
     "rz": braket_gates.Rz,
+    "U": braket_gates.U,
     "swap": braket_gates.Swap,
     "iswap": braket_gates.ISwap,
     "pswap": braket_gates.PSwap,

test/unit_tests/braket/circuits/test_gates.py

@@ -54,6 +54,7 @@ class TripleAngle:
     (Gate.Rx, "rx", ir.Rx, [SingleTarget, Angle], {}),
     (Gate.Ry, "ry", ir.Ry, [SingleTarget, Angle], {}),
     (Gate.Rz, "rz", ir.Rz, [SingleTarget, Angle], {}),
+    (Gate.U, "u", None, [SingleTarget, TripleAngle], {}),
     (Gate.CNot, "cnot", ir.CNot, [SingleTarget, SingleControl], {}),
     (Gate.CV, "cv", ir.CV, [SingleTarget, SingleControl], {}),
     (Gate.CCNot, "ccnot", ir.CCNot, [SingleTarget, DoubleControl], {}),
@@ -121,6 +122,7 @@ class TripleAngle:
     Gate.Rx,
     Gate.Ry,
     Gate.Rz,
+    Gate.U,
     Gate.PhaseShift,
     Gate.PSwap,
     Gate.XX,
@@ -434,6 +436,18 @@ def test_ir_gate_level(testclass, subroutine_name, irclass, irsubclasses, kwargs
             OpenQASMSerializationProperties(qubit_reference_type=QubitReferenceType.PHYSICAL),
             "rz(0.17) $4;",
         ),
+        (
+            Gate.U(angle_1=0.17, angle_2=3.45, angle_3=5.21),
+            [4],
+            OpenQASMSerializationProperties(qubit_reference_type=QubitReferenceType.VIRTUAL),
+            "U(0.17, 3.45, 5.21) q[4];",
+        ),
+        (
+            Gate.U(angle_1=0.17, angle_2=3.45, angle_3=5.21),
+            [4],
+            OpenQASMSerializationProperties(qubit_reference_type=QubitReferenceType.PHYSICAL),
+            "U(0.17, 3.45, 5.21) $4;",
+        ),
         (
             Gate.XX(angle=0.17),
             [4, 5],
@@ -859,6 +873,8 @@ def test_gate_subroutine(testclass, subroutine_name, irclass, irsubclasses, kwar
         subroutine_input = {"target": multi_targets}
         if Angle in irsubclasses:
             subroutine_input.update(angle_valid_input())
+        if TripleAngle in irsubclasses:
+            subroutine_input.update(triple_angle_valid_input())
         assert subroutine(**subroutine_input) == Circuit(instruction_list)
 
 
@@ -925,7 +941,7 @@ def test_large_unitary():
 
 @pytest.mark.parametrize("gate", parameterizable_gates)
 def test_bind_values(gate):
-    triple_angled = gate.__name__ in ("MS",)
+    triple_angled = gate.__name__ in ("MS", "U")
     num_params = 3 if triple_angled else 1
     thetas = [FreeParameter(f"theta_{i}") for i in range(num_params)]
     mapping = {f"theta_{i}": i for i in range(num_params)}