Align as parametrized gate

src/qibo/backends/npmatrices.py

@@ -64,7 +64,7 @@ def I(self, n=2):
         # _cast will take care of casting in the right dtype for all the backends
         return self._cast(self.np.eye(n, dtype=complex), dtype=self.dtype)
 
-    def Align(self, n=2):
+    def Align(self, delay, n=2):
         return self._cast(self.I(n), dtype=self.dtype)
 
     def M(self):  # pragma: no cover

src/qibo/gates/gates.py

@@ -503,30 +503,31 @@ def qasm_label(self):
         return "id"
 
 
-class Align(Gate):
+class Align(ParametrizedGate):
     """Aligns proceeding qubit operations and (optionally) waits ``delay`` amount of time.
 
     Args:
-        *q (int): The qubit ID numbers.
+        q (int): The qubit ID.
         delay (int, optional): The time (in ns) for which to delay circuit execution on the specified qubits.
             Defaults to ``0`` (zero).
     """
 
-    def __init__(self, *q, delay: int = 0):
+    def __init__(self, q, delay=0, trainable=True):
         if not isinstance(delay, int):
             raise_error(
                 TypeError, f"delay must be type int, but it is type {type(delay)}."
             )
         if delay < 0.0:
             raise_error(ValueError, "Delay must not be negative.")
 
-        super().__init__()
+        super().__init__(trainable)
         self.name = "align"
-        self.delay = delay
         self.draw_label = f"A({delay})"
-        self.init_args = q
-        self.init_kwargs = {"delay": delay}
-        self.target_qubits = tuple(q)
+        self.init_args = [q]
+        self.init_kwargs = {"name": self.name, "delay": delay, "trainable": trainable}
+        self.target_qubits = (q,)
+        self._parameters = (delay,)
+        self.nparams = 1
 
 
 def _is_clifford_given_angle(angle):

tests/test_gates_gates.py

@@ -212,12 +212,12 @@ def test_align(backend):
     with pytest.raises(ValueError):
         gates.Align(0, delay=-1)
 
-    nqubits = 2
+    nqubits = 1
 
-    gate = gates.Align(0, 1)
-    gate_list = [gates.H(0), gates.H(1), gate]
+    gate = gates.Align(0, 0)
+    gate_list = [gates.H(0), gate]
 
-    final_state = apply_gates(backend, gate_list, nqubits=2)
+    final_state = apply_gates(backend, gate_list, nqubits=nqubits)
     target_state = backend.plus_state(nqubits)
 
     backend.assert_allclose(final_state, target_state)
@@ -229,8 +229,8 @@ def test_align(backend):
     with pytest.raises(NotImplementedError):
         gate.qasm_label
 
-    assert not gates.Align(0, 1).clifford
-    assert not gates.Align(0, 1).unitary
+    assert not gates.Align(0, delay=0).clifford
+    assert not gates.Align(0, delay=0).unitary
 
 
 # :class:`qibo.core.cgates.M` is tested seperately in `test_measurement_gate.py`