Repository: qiboteam/qibo

Issue:
Align gate should be a `ParametrizedGate`
Currently the `Align` gate accepts as parameter `delay` however it is not treated as a "parameter" of the gate since in the current implementation `Align` is defined from `Gate` and not from `ParametrizedGate`.
I think we could define it directly as a `ParametrizedGate` 

Discussion:

Patch:
diff --git a/src/qibo/backends/npmatrices.py b/src/qibo/backends/npmatrices.py
index c6e447e9e..184ef9c54 100644
--- a/src/qibo/backends/npmatrices.py
+++ b/src/qibo/backends/npmatrices.py
@@ -64,7 +64,7 @@ class NumpyMatrices:
         # _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
diff --git a/src/qibo/gates/gates.py b/src/qibo/gates/gates.py
index dfbe628a0..435ce821a 100644
--- a/src/qibo/gates/gates.py
+++ b/src/qibo/gates/gates.py
@@ -503,16 +503,16 @@ class I(Gate):
         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)}."
@@ -520,13 +520,14 @@ class Align(Gate):
         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):


Test patch:
diff --git a/tests/test_gates_gates.py b/tests/test_gates_gates.py
index d407e72e9..dbd04d9be 100644
--- a/tests/test_gates_gates.py
+++ b/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`