From e468f96a25b2bc185efd93ef0e4af74a19ea8ed0 Mon Sep 17 00:00:00 2001
From: EthanObadia <ethan.obadia@pasqal.com>
Date: Fri, 5 Apr 2024 15:51:30 +0200
Subject: [PATCH 01/26] add promote_ope function (#154)

---
 pyqtorch/utils.py | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/pyqtorch/utils.py b/pyqtorch/utils.py
index 43cbd3e7..a7b3e4ad 100644
--- a/pyqtorch/utils.py
+++ b/pyqtorch/utils.py
@@ -7,6 +7,7 @@
 from torch import Tensor
 
 from pyqtorch.matrices import DEFAULT_MATRIX_DTYPE, DEFAULT_REAL_DTYPE
+from pyqtorch.primitive import I
 
 State = Tensor
 Operator = Tensor
@@ -142,3 +143,31 @@ def density_mat(state: Tensor) -> Tensor:
     state = torch.permute(state, batch_first_perm).reshape(batch_size, 2**n_qubits)
     undo_perm = (1, 2, 0)
     return torch.permute(torch.einsum("bi,bj->bij", (state, state.conj())), undo_perm)
+
+
+def promote_ope(operator: Tensor, target: int, n_qubits: int) -> Tensor:
+    """
+    Promotes `operator` to the size of the circuit.
+
+    Args:
+        operator (Tensor): The operator tensor to be promoted.
+        target (int): The target qubit index.
+        n_qubits (int): Number of qubits in the circuit.
+
+    Returns:
+        Tensor: The promoted operator tensor.
+
+    Raises:
+        ValueError: If `target` is not within the range of qubits.
+    """
+    if target > n_qubits:
+        raise ValueError("The target must be a register qubit")
+
+    qubits_support = torch.arange(0, n_qubits)
+    for support in qubits_support:
+        if target > support:
+            operator = torch.kron(I(support).unitary(), operator.contiguous())
+        # Add .contiguous() because kron does not support the transpose (dagger)
+        elif target < support:
+            operator = torch.kron(operator.contiguous(), I(support).unitary())
+    return operator
\ No newline at end of file

From aa0ca93614b930effa042a5a0cef7604e0d71c7d Mon Sep 17 00:00:00 2001
From: EthanObadia <ethan.obadia@pasqal.com>
Date: Fri, 5 Apr 2024 16:08:03 +0200
Subject: [PATCH 02/26] add apply_ope_ope() function (#153)

---
 pyqtorch/apply.py | 67 +++++++++++++++++++++++++++++++++++++++++++++--
 pyqtorch/utils.py | 29 ++++++++++++++++++++
 2 files changed, 94 insertions(+), 2 deletions(-)

diff --git a/pyqtorch/apply.py b/pyqtorch/apply.py
index 859f7b82..b5082e32 100644
--- a/pyqtorch/apply.py
+++ b/pyqtorch/apply.py
@@ -3,11 +3,14 @@
 from string import ascii_letters as ABC
 from typing import Tuple
 
+from math import log2
 from numpy import array
 from numpy.typing import NDArray
-from torch import einsum
 
-from pyqtorch.utils import Operator, State
+import torch
+from torch import Tensor,einsum
+
+from pyqtorch.utils import Operator, State, promote_ope
 
 ABC_ARRAY: NDArray = array(list(ABC))
 
@@ -57,3 +60,63 @@ def apply_operator(
         map(lambda e: "".join(list(e)), [operator_dims, in_state_dims, out_state_dims])
     )
     return einsum(f"{operator_dims},{in_state_dims}->{out_state_dims}", operator, state)
+
+
+def apply_ope_ope(operator_1: Tensor, operator_2: Tensor, target: int) -> Tensor:
+    """
+    Compute the product of two operators.
+
+    Args:
+        operator_1 (Tensor): The first operator.
+        operator_2 (Tensor): The second operator.
+        target (int): The target qubit index.
+
+    Returns:
+        Tensor: The product of the two operators.
+    """
+
+    # Dimension verifications:
+    n_qubits_1 = int(log2(operator_1.size(1)))
+    n_qubits_2 = int(log2(operator_2.size(1)))
+    batch_size_1 = operator_1.size(-1)
+    batch_size_2 = operator_2.size(-1)
+    if n_qubits_1 != n_qubits_2:
+        if n_qubits_1 > n_qubits_2:
+            operator_2 = promote_ope(operator_2, target, n_qubits_1)
+        if n_qubits_1 < n_qubits_2:
+            operator_1 = promote_ope(operator_1, target, n_qubits_2)
+    if batch_size_1 != batch_size_2:
+        if batch_size_1 > batch_size_2:
+            operator_2 = operator_2.repeat(1, 1, batch_size_1)
+        if batch_size_2 > batch_size_1:
+            operator_1 = operator_1.repeat(1, 1, batch_size_2)
+
+    # Permute the batch size on first dimension to allow torch.bmm():
+    def batch_first(operator: Tensor) -> Tensor:
+        """
+        Permute the operator's batch dimension on first dimension.
+
+        Args:
+        operator (Tensor): Operator in size [2**n_qubits, 2**n_qubits,batch_size].
+
+        Returns:
+        Tensor: Operator in size [batch_size, 2**n_qubits, 2**n_qubits].
+        """
+        batch_first_perm = (2, 0, 1)
+        return torch.permute(operator, batch_first_perm)
+
+    # Undo the permute since PyQ expects tensor.Size([2**n_qubits, 2**n_qubits,batch_size]):
+    def batch_last(operator: Tensor) -> Tensor:
+        """
+        Permute the operator's batch dimension on last dimension.
+
+        Args:
+        operator (Tensor): Operator in size [batch_size,2**n_qubits, 2**n_qubits].
+
+        Returns:
+        Tensor: Operator in size [2**n_qubits, 2**n_qubits,batch_size].
+        """
+        undo_perm = (1, 2, 0)
+        return torch.permute(operator, undo_perm)
+
+    return batch_last(torch.bmm(batch_first(operator_1), batch_first(operator_2)))
\ No newline at end of file
diff --git a/pyqtorch/utils.py b/pyqtorch/utils.py
index 43cbd3e7..a7b3e4ad 100644
--- a/pyqtorch/utils.py
+++ b/pyqtorch/utils.py
@@ -7,6 +7,7 @@
 from torch import Tensor
 
 from pyqtorch.matrices import DEFAULT_MATRIX_DTYPE, DEFAULT_REAL_DTYPE
+from pyqtorch.primitive import I
 
 State = Tensor
 Operator = Tensor
@@ -142,3 +143,31 @@ def density_mat(state: Tensor) -> Tensor:
     state = torch.permute(state, batch_first_perm).reshape(batch_size, 2**n_qubits)
     undo_perm = (1, 2, 0)
     return torch.permute(torch.einsum("bi,bj->bij", (state, state.conj())), undo_perm)
+
+
+def promote_ope(operator: Tensor, target: int, n_qubits: int) -> Tensor:
+    """
+    Promotes `operator` to the size of the circuit.
+
+    Args:
+        operator (Tensor): The operator tensor to be promoted.
+        target (int): The target qubit index.
+        n_qubits (int): Number of qubits in the circuit.
+
+    Returns:
+        Tensor: The promoted operator tensor.
+
+    Raises:
+        ValueError: If `target` is not within the range of qubits.
+    """
+    if target > n_qubits:
+        raise ValueError("The target must be a register qubit")
+
+    qubits_support = torch.arange(0, n_qubits)
+    for support in qubits_support:
+        if target > support:
+            operator = torch.kron(I(support).unitary(), operator.contiguous())
+        # Add .contiguous() because kron does not support the transpose (dagger)
+        elif target < support:
+            operator = torch.kron(operator.contiguous(), I(support).unitary())
+    return operator
\ No newline at end of file

From 6606021c0ec812f5a66c66e37081d585573787de Mon Sep 17 00:00:00 2001
From: EthanObadia <ethan.obadia@pasqal.com>
Date: Fri, 5 Apr 2024 17:27:55 +0200
Subject: [PATCH 03/26] Modify the target range value (#154)

---
 pyqtorch/utils.py | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/pyqtorch/utils.py b/pyqtorch/utils.py
index a7b3e4ad..a9e0a3a7 100644
--- a/pyqtorch/utils.py
+++ b/pyqtorch/utils.py
@@ -147,27 +147,29 @@ def density_mat(state: Tensor) -> Tensor:
 
 def promote_ope(operator: Tensor, target: int, n_qubits: int) -> Tensor:
     """
-    Promotes `operator` to the size of the circuit.
+    Promotes `operator` to the size of the circuit (number of qubits and batch).
+    Targeting the first qubit implies target = 0, so target > n_qubits - 1.
 
     Args:
         operator (Tensor): The operator tensor to be promoted.
-        target (int): The target qubit index.
+        target (int): The index of the target qubit to which the operator is applied.
+            Targeting the first qubit implies target = 0, so target > n_qubits - 1.
         n_qubits (int): Number of qubits in the circuit.
 
     Returns:
         Tensor: The promoted operator tensor.
 
     Raises:
-        ValueError: If `target` is not within the range of qubits.
+        ValueError: If `target` is outside the valid range of qubits.
     """
-    if target > n_qubits:
-        raise ValueError("The target must be a register qubit")
+    if target > n_qubits - 1 : 
+        raise ValueError("The target must be a valid qubit index within the circuit's range.")
 
     qubits_support = torch.arange(0, n_qubits)
     for support in qubits_support:
         if target > support:
             operator = torch.kron(I(support).unitary(), operator.contiguous())
-        # Add .contiguous() because kron does not support the transpose (dagger)
+        # Add.contiguous() because kron does not support the transpose (dagger)
         elif target < support:
             operator = torch.kron(operator.contiguous(), I(support).unitary())
     return operator
\ No newline at end of file

From 639bbdb1d69ce98132baf1395636511bd8f4d489 Mon Sep 17 00:00:00 2001
From: EthanObadia <ethan.obadia@pasqal.com>
Date: Tue, 9 Apr 2024 10:33:51 +0200
Subject: [PATCH 04/26] Add test_promote func (#155)

---
 pyqtorch/apply.py     | 10 ++++------
 pyqtorch/primitive.py | 12 +++++++-----
 pyqtorch/utils.py     |  4 ++--
 tests/test_digital.py | 13 ++++++++++++-
 4 files changed, 25 insertions(+), 14 deletions(-)

diff --git a/pyqtorch/apply.py b/pyqtorch/apply.py
index 859f7b82..f7f5b151 100644
--- a/pyqtorch/apply.py
+++ b/pyqtorch/apply.py
@@ -5,20 +5,18 @@
 
 from numpy import array
 from numpy.typing import NDArray
-from torch import einsum
-
-from pyqtorch.utils import Operator, State
+from torch import Tensor, einsum
 
 ABC_ARRAY: NDArray = array(list(ABC))
 
 
 def apply_operator(
-    state: State,
-    operator: Operator,
+    state: Tensor,
+    operator: Tensor,
     qubits: Tuple[int, ...] | list[int],
     n_qubits: int = None,
     batch_size: int = None,
-) -> State:
+) -> Tensor:
     """Applies an operator, i.e. a single tensor of shape [2, 2, ...], on a given state
        of shape [2 for _ in range(n_qubits)] for a given set of (target and control) qubits.
 
diff --git a/pyqtorch/primitive.py b/pyqtorch/primitive.py
index bac9e4ba..2af104bb 100644
--- a/pyqtorch/primitive.py
+++ b/pyqtorch/primitive.py
@@ -7,7 +7,7 @@
 
 from pyqtorch.apply import apply_operator
 from pyqtorch.matrices import OPERATIONS_DICT, _controlled, _dagger
-from pyqtorch.utils import Operator, State, product_state
+from pyqtorch.utils import product_state
 
 
 class Primitive(torch.nn.Module):
@@ -40,15 +40,17 @@ def extra_repr(self) -> str:
     def param_type(self) -> None:
         return self._param_type
 
-    def unitary(self, values: dict[str, torch.Tensor] | torch.Tensor = {}) -> Operator:
+    def unitary(self, values: dict[str, torch.Tensor] | torch.Tensor = {}) -> torch.Tensor:
         return self.pauli.unsqueeze(2)
 
-    def forward(self, state: State, values: dict[str, torch.Tensor] | torch.Tensor = {}) -> State:
+    def forward(
+        self, state: torch.Tensor, values: dict[str, torch.Tensor] | torch.Tensor = {}
+    ) -> torch.Tensor:
         return apply_operator(
             state, self.unitary(values), self.qubit_support, len(state.size()) - 1
         )
 
-    def dagger(self, values: dict[str, torch.Tensor] | torch.Tensor = {}) -> Operator:
+    def dagger(self, values: dict[str, torch.Tensor] | torch.Tensor = {}) -> torch.Tensor:
         return _dagger(self.unitary(values))
 
     @property
@@ -85,7 +87,7 @@ class I(Primitive):  # noqa: E742
     def __init__(self, target: int):
         super().__init__(OPERATIONS_DICT["I"], target)
 
-    def forward(self, state: State, values: dict[str, torch.Tensor] = None) -> State:
+    def forward(self, state: torch.Tensor, values: dict[str, torch.Tensor] = None) -> torch.Tensor:
         return state
 
 
diff --git a/pyqtorch/utils.py b/pyqtorch/utils.py
index a9e0a3a7..40c72ec7 100644
--- a/pyqtorch/utils.py
+++ b/pyqtorch/utils.py
@@ -162,7 +162,7 @@ def promote_ope(operator: Tensor, target: int, n_qubits: int) -> Tensor:
     Raises:
         ValueError: If `target` is outside the valid range of qubits.
     """
-    if target > n_qubits - 1 : 
+    if target > n_qubits - 1:
         raise ValueError("The target must be a valid qubit index within the circuit's range.")
 
     qubits_support = torch.arange(0, n_qubits)
@@ -172,4 +172,4 @@ def promote_ope(operator: Tensor, target: int, n_qubits: int) -> Tensor:
         # Add.contiguous() because kron does not support the transpose (dagger)
         elif target < support:
             operator = torch.kron(operator.contiguous(), I(support).unitary())
-    return operator
\ No newline at end of file
+    return operator
diff --git a/tests/test_digital.py b/tests/test_digital.py
index d67b470a..6c6cdfc0 100644
--- a/tests/test_digital.py
+++ b/tests/test_digital.py
@@ -12,7 +12,8 @@
 from pyqtorch.apply import apply_operator
 from pyqtorch.matrices import DEFAULT_MATRIX_DTYPE, IMAT, ZMAT
 from pyqtorch.parametric import Parametric
-from pyqtorch.utils import ATOL, density_mat, product_state, random_state
+from pyqtorch.primitive import I
+from pyqtorch.utils import ATOL, density_mat, product_state, promote_ope, random_state
 
 state_000 = product_state("000")
 state_001 = product_state("001")
@@ -322,3 +323,13 @@ def test_dm(n_qubits: Tensor, batch_size: Tensor) -> None:
     dm = density_mat(state_cat)
     assert dm.size() == torch.Size([2**n_qubits, 2**n_qubits, batch_size])
     assert torch.allclose(dm, dm_proj)
+
+
+random_number = torch.randint(1, 6, (8, 2))
+sorted_numbers, _ = torch.sort(random_number, dim=1)
+
+
+@pytest.mark.parametrize("target,n_qubits", sorted_numbers)
+def test_promote(target: int, n_qubits: int) -> None:
+    operator: Tensor = promote_ope(I(target).unitary(), target, n_qubits)
+    assert operator.size() == torch.Size([2**n_qubits, 2**n_qubits, 1])

From 2bac1b01af80e96b0be96560fc32dd7e912a3ef5 Mon Sep 17 00:00:00 2001
From: EthanObadia <ethan.obadia@pasqal.com>
Date: Tue, 9 Apr 2024 10:57:06 +0200
Subject: [PATCH 05/26] Modify test_promote func (#155)

---
 tests/test_digital.py | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/tests/test_digital.py b/tests/test_digital.py
index 6c6cdfc0..80bf8512 100644
--- a/tests/test_digital.py
+++ b/tests/test_digital.py
@@ -12,7 +12,7 @@
 from pyqtorch.apply import apply_operator
 from pyqtorch.matrices import DEFAULT_MATRIX_DTYPE, IMAT, ZMAT
 from pyqtorch.parametric import Parametric
-from pyqtorch.primitive import I
+from pyqtorch.primitive import I, X
 from pyqtorch.utils import ATOL, density_mat, product_state, promote_ope, random_state
 
 state_000 = product_state("000")
@@ -325,11 +325,15 @@ def test_dm(n_qubits: Tensor, batch_size: Tensor) -> None:
     assert torch.allclose(dm, dm_proj)
 
 
-random_number = torch.randint(1, 6, (8, 2))
-sorted_numbers, _ = torch.sort(random_number, dim=1)
+size = (5, 2)
+random_param = torch.randperm(size[0] * size[1])
+random_param = random_param.view(size)
+random_param = torch.sort(random_param, dim=1)[0]
 
 
-@pytest.mark.parametrize("target,n_qubits", sorted_numbers)
+@pytest.mark.parametrize("target,n_qubits", random_param)
 def test_promote(target: int, n_qubits: int) -> None:
-    operator: Tensor = promote_ope(I(target).unitary(), target, n_qubits)
-    assert operator.size() == torch.Size([2**n_qubits, 2**n_qubits, 1])
+    I_prom = promote_ope(I(0).unitary(), target, n_qubits)
+    assert I_prom.size() == torch.Size([2**n_qubits, 2**n_qubits, 1])
+    X_prom = promote_ope(X(0).unitary(), target, n_qubits)
+    assert X_prom.size() == torch.Size([2**n_qubits, 2**n_qubits, 1])

From 047a3c20530731306531d87b581784c9d20d736f Mon Sep 17 00:00:00 2001
From: EthanObadia <ethan.obadia@pasqal.com>
Date: Mon, 15 Apr 2024 12:24:14 +0200
Subject: [PATCH 06/26] Move batch first/last functions in utils (#153)

---
 pyqtorch/apply.py | 30 +-----------------------------
 pyqtorch/utils.py | 30 +++++++++++++++++++++++++++++-
 2 files changed, 30 insertions(+), 30 deletions(-)

diff --git a/pyqtorch/apply.py b/pyqtorch/apply.py
index b5082e32..52b2fc28 100644
--- a/pyqtorch/apply.py
+++ b/pyqtorch/apply.py
@@ -10,7 +10,7 @@
 import torch
 from torch import Tensor,einsum
 
-from pyqtorch.utils import Operator, State, promote_ope
+from pyqtorch.utils import Operator, State, promote_ope,batch_first, batch_last
 
 ABC_ARRAY: NDArray = array(list(ABC))
 
@@ -91,32 +91,4 @@ def apply_ope_ope(operator_1: Tensor, operator_2: Tensor, target: int) -> Tensor
         if batch_size_2 > batch_size_1:
             operator_1 = operator_1.repeat(1, 1, batch_size_2)
 
-    # Permute the batch size on first dimension to allow torch.bmm():
-    def batch_first(operator: Tensor) -> Tensor:
-        """
-        Permute the operator's batch dimension on first dimension.
-
-        Args:
-        operator (Tensor): Operator in size [2**n_qubits, 2**n_qubits,batch_size].
-
-        Returns:
-        Tensor: Operator in size [batch_size, 2**n_qubits, 2**n_qubits].
-        """
-        batch_first_perm = (2, 0, 1)
-        return torch.permute(operator, batch_first_perm)
-
-    # Undo the permute since PyQ expects tensor.Size([2**n_qubits, 2**n_qubits,batch_size]):
-    def batch_last(operator: Tensor) -> Tensor:
-        """
-        Permute the operator's batch dimension on last dimension.
-
-        Args:
-        operator (Tensor): Operator in size [batch_size,2**n_qubits, 2**n_qubits].
-
-        Returns:
-        Tensor: Operator in size [2**n_qubits, 2**n_qubits,batch_size].
-        """
-        undo_perm = (1, 2, 0)
-        return torch.permute(operator, undo_perm)
-
     return batch_last(torch.bmm(batch_first(operator_1), batch_first(operator_2)))
\ No newline at end of file
diff --git a/pyqtorch/utils.py b/pyqtorch/utils.py
index a7b3e4ad..3d6e353c 100644
--- a/pyqtorch/utils.py
+++ b/pyqtorch/utils.py
@@ -170,4 +170,32 @@ def promote_ope(operator: Tensor, target: int, n_qubits: int) -> Tensor:
         # Add .contiguous() because kron does not support the transpose (dagger)
         elif target < support:
             operator = torch.kron(operator.contiguous(), I(support).unitary())
-    return operator
\ No newline at end of file
+    return operator
+
+
+def batch_first(operator: Tensor) -> Tensor:
+    """
+    Permute the operator's batch dimension on first dimension.
+
+    Args:
+    operator (Tensor): Operator in size [2**n_qubits, 2**n_qubits,batch_size].
+
+    Returns:
+    Tensor: Operator in size [batch_size, 2**n_qubits, 2**n_qubits].
+    """
+    batch_first_perm = (2, 0, 1)
+    return torch.permute(operator, batch_first_perm)
+
+
+def batch_last(operator: Tensor) -> Tensor:
+    """
+    Permute the operator's batch dimension on last dimension.
+
+    Args:
+    operator (Tensor): Operator in size [batch_size,2**n_qubits, 2**n_qubits].
+
+    Returns:
+    Tensor: Operator in size [2**n_qubits, 2**n_qubits,batch_size].
+    """
+    undo_perm = (1, 2, 0)
+    return torch.permute(operator, undo_perm)
\ No newline at end of file

From 817fb0e72c37703496f5e90b1aba03dd3aa2da9c Mon Sep 17 00:00:00 2001
From: EthanObadia <ethan.obadia@pasqal.com>
Date: Mon, 15 Apr 2024 12:37:31 +0200
Subject: [PATCH 07/26] pre-commit

---
 pyqtorch/apply.py     | 17 ++++++++---------
 pyqtorch/primitive.py | 11 ++++++-----
 pyqtorch/utils.py     |  2 +-
 3 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/pyqtorch/apply.py b/pyqtorch/apply.py
index 52b2fc28..e0748df6 100644
--- a/pyqtorch/apply.py
+++ b/pyqtorch/apply.py
@@ -1,27 +1,26 @@
 from __future__ import annotations
 
+from math import log2
 from string import ascii_letters as ABC
 from typing import Tuple
 
-from math import log2
+import torch
 from numpy import array
 from numpy.typing import NDArray
+from torch import Tensor, einsum
 
-import torch
-from torch import Tensor,einsum
-
-from pyqtorch.utils import Operator, State, promote_ope,batch_first, batch_last
+from pyqtorch.utils import batch_first, batch_last, promote_ope
 
 ABC_ARRAY: NDArray = array(list(ABC))
 
 
 def apply_operator(
-    state: State,
-    operator: Operator,
+    state: Tensor,
+    operator: Tensor,
     qubits: Tuple[int, ...] | list[int],
     n_qubits: int = None,
     batch_size: int = None,
-) -> State:
+) -> Tensor:
     """Applies an operator, i.e. a single tensor of shape [2, 2, ...], on a given state
        of shape [2 for _ in range(n_qubits)] for a given set of (target and control) qubits.
 
@@ -91,4 +90,4 @@ def apply_ope_ope(operator_1: Tensor, operator_2: Tensor, target: int) -> Tensor
         if batch_size_2 > batch_size_1:
             operator_1 = operator_1.repeat(1, 1, batch_size_2)
 
-    return batch_last(torch.bmm(batch_first(operator_1), batch_first(operator_2)))
\ No newline at end of file
+    return batch_last(torch.bmm(batch_first(operator_1), batch_first(operator_2)))
diff --git a/pyqtorch/primitive.py b/pyqtorch/primitive.py
index bac9e4ba..66dcfb27 100644
--- a/pyqtorch/primitive.py
+++ b/pyqtorch/primitive.py
@@ -4,10 +4,11 @@
 from typing import Any, Tuple
 
 import torch
+from torch import Tensor
 
 from pyqtorch.apply import apply_operator
 from pyqtorch.matrices import OPERATIONS_DICT, _controlled, _dagger
-from pyqtorch.utils import Operator, State, product_state
+from pyqtorch.utils import product_state
 
 
 class Primitive(torch.nn.Module):
@@ -40,15 +41,15 @@ def extra_repr(self) -> str:
     def param_type(self) -> None:
         return self._param_type
 
-    def unitary(self, values: dict[str, torch.Tensor] | torch.Tensor = {}) -> Operator:
+    def unitary(self, values: dict[str, torch.Tensor] | torch.Tensor = {}) -> Tensor:
         return self.pauli.unsqueeze(2)
 
-    def forward(self, state: State, values: dict[str, torch.Tensor] | torch.Tensor = {}) -> State:
+    def forward(self, state: Tensor, values: dict[str, torch.Tensor] | torch.Tensor = {}) -> Tensor:
         return apply_operator(
             state, self.unitary(values), self.qubit_support, len(state.size()) - 1
         )
 
-    def dagger(self, values: dict[str, torch.Tensor] | torch.Tensor = {}) -> Operator:
+    def dagger(self, values: dict[str, torch.Tensor] | torch.Tensor = {}) -> Tensor:
         return _dagger(self.unitary(values))
 
     @property
@@ -85,7 +86,7 @@ class I(Primitive):  # noqa: E742
     def __init__(self, target: int):
         super().__init__(OPERATIONS_DICT["I"], target)
 
-    def forward(self, state: State, values: dict[str, torch.Tensor] = None) -> State:
+    def forward(self, state: Tensor, values: dict[str, torch.Tensor] = None) -> Tensor:
         return state
 
 
diff --git a/pyqtorch/utils.py b/pyqtorch/utils.py
index 3d6e353c..bf20fc99 100644
--- a/pyqtorch/utils.py
+++ b/pyqtorch/utils.py
@@ -198,4 +198,4 @@ def batch_last(operator: Tensor) -> Tensor:
     Tensor: Operator in size [2**n_qubits, 2**n_qubits,batch_size].
     """
     undo_perm = (1, 2, 0)
-    return torch.permute(operator, undo_perm)
\ No newline at end of file
+    return torch.permute(operator, undo_perm)

From ab95bb9417b1c186925ac771dd6dde08d4a54bc8 Mon Sep 17 00:00:00 2001
From: EthanObadia <ethan.obadia@pasqal.com>
Date: Mon, 15 Apr 2024 12:57:10 +0200
Subject: [PATCH 08/26] Modify variable names (#154)

---
 pyqtorch/utils.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/pyqtorch/utils.py b/pyqtorch/utils.py
index 40c72ec7..159c8aa2 100644
--- a/pyqtorch/utils.py
+++ b/pyqtorch/utils.py
@@ -165,11 +165,11 @@ def promote_ope(operator: Tensor, target: int, n_qubits: int) -> Tensor:
     if target > n_qubits - 1:
         raise ValueError("The target must be a valid qubit index within the circuit's range.")
 
-    qubits_support = torch.arange(0, n_qubits)
-    for support in qubits_support:
-        if target > support:
-            operator = torch.kron(I(support).unitary(), operator.contiguous())
-        # Add.contiguous() because kron does not support the transpose (dagger)
-        elif target < support:
-            operator = torch.kron(operator.contiguous(), I(support).unitary())
+    qubits = torch.arange(0, n_qubits)
+    for qubit in qubits:
+        if target > qubit:
+            operator = torch.kron(I(qubit).unitary(), operator.contiguous())
+        # Add .contiguous() because kron does not support the transpose (dagger)
+        elif target < qubit:
+            operator = torch.kron(operator.contiguous(), I(qubit).unitary())
     return operator

From b1ce1d7b28fb3e5128200dbe86f2420049766f95 Mon Sep 17 00:00:00 2001
From: EthanObadia <ethan.obadia@pasqal.com>
Date: Mon, 15 Apr 2024 17:49:27 +0200
Subject: [PATCH 09/26] Modify after review (#155)

---
 pyqtorch/primitive.py | 13 ++++++-------
 pyqtorch/utils.py     |  9 +++++----
 tests/test_digital.py |  6 +++---
 3 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/pyqtorch/primitive.py b/pyqtorch/primitive.py
index 2af104bb..f6683377 100644
--- a/pyqtorch/primitive.py
+++ b/pyqtorch/primitive.py
@@ -4,6 +4,7 @@
 from typing import Any, Tuple
 
 import torch
+from torch import Tensor
 
 from pyqtorch.apply import apply_operator
 from pyqtorch.matrices import OPERATIONS_DICT, _controlled, _dagger
@@ -11,7 +12,7 @@
 
 
 class Primitive(torch.nn.Module):
-    def __init__(self, pauli: torch.Tensor, target: int) -> None:
+    def __init__(self, pauli: Tensor, target: int) -> None:
         super().__init__()
         self.target: int = target
         self.qubit_support: Tuple[int, ...] = (target,)
@@ -40,17 +41,15 @@ def extra_repr(self) -> str:
     def param_type(self) -> None:
         return self._param_type
 
-    def unitary(self, values: dict[str, torch.Tensor] | torch.Tensor = {}) -> torch.Tensor:
+    def unitary(self, values: dict[str, Tensor] | Tensor = {}) -> Tensor:
         return self.pauli.unsqueeze(2)
 
-    def forward(
-        self, state: torch.Tensor, values: dict[str, torch.Tensor] | torch.Tensor = {}
-    ) -> torch.Tensor:
+    def forward(self, state: Tensor, values: dict[str, Tensor] | Tensor = {}) -> Tensor:
         return apply_operator(
             state, self.unitary(values), self.qubit_support, len(state.size()) - 1
         )
 
-    def dagger(self, values: dict[str, torch.Tensor] | torch.Tensor = {}) -> torch.Tensor:
+    def dagger(self, values: dict[str, Tensor] | Tensor = {}) -> Tensor:
         return _dagger(self.unitary(values))
 
     @property
@@ -87,7 +86,7 @@ class I(Primitive):  # noqa: E742
     def __init__(self, target: int):
         super().__init__(OPERATIONS_DICT["I"], target)
 
-    def forward(self, state: torch.Tensor, values: dict[str, torch.Tensor] = None) -> torch.Tensor:
+    def forward(self, state: Tensor, values: dict[str, Tensor] = None) -> Tensor:
         return state
 
 
diff --git a/pyqtorch/utils.py b/pyqtorch/utils.py
index 159c8aa2..b91b44a1 100644
--- a/pyqtorch/utils.py
+++ b/pyqtorch/utils.py
@@ -7,7 +7,6 @@
 from torch import Tensor
 
 from pyqtorch.matrices import DEFAULT_MATRIX_DTYPE, DEFAULT_REAL_DTYPE
-from pyqtorch.primitive import I
 
 State = Tensor
 Operator = Tensor
@@ -145,7 +144,7 @@ def density_mat(state: Tensor) -> Tensor:
     return torch.permute(torch.einsum("bi,bj->bij", (state, state.conj())), undo_perm)
 
 
-def promote_ope(operator: Tensor, target: int, n_qubits: int) -> Tensor:
+def promote_op(operator: Tensor, target: int, n_qubits: int) -> Tensor:
     """
     Promotes `operator` to the size of the circuit (number of qubits and batch).
     Targeting the first qubit implies target = 0, so target > n_qubits - 1.
@@ -166,10 +165,12 @@ def promote_ope(operator: Tensor, target: int, n_qubits: int) -> Tensor:
         raise ValueError("The target must be a valid qubit index within the circuit's range.")
 
     qubits = torch.arange(0, n_qubits)
+    # Define I instead of importing it from pyqtorch.primitive to avoid circular import
+    identity = torch.tensor([[[1.0 + 0.0j], [0.0 + 0.0j]], [[0.0 + 0.0j], [1.0 + 0.0j]]])
     for qubit in qubits:
         if target > qubit:
-            operator = torch.kron(I(qubit).unitary(), operator.contiguous())
+            operator = torch.kron(identity, operator.contiguous())
         # Add .contiguous() because kron does not support the transpose (dagger)
         elif target < qubit:
-            operator = torch.kron(operator.contiguous(), I(qubit).unitary())
+            operator = torch.kron(operator.contiguous(), identity)
     return operator
diff --git a/tests/test_digital.py b/tests/test_digital.py
index 80bf8512..ad5d0661 100644
--- a/tests/test_digital.py
+++ b/tests/test_digital.py
@@ -13,7 +13,7 @@
 from pyqtorch.matrices import DEFAULT_MATRIX_DTYPE, IMAT, ZMAT
 from pyqtorch.parametric import Parametric
 from pyqtorch.primitive import I, X
-from pyqtorch.utils import ATOL, density_mat, product_state, promote_ope, random_state
+from pyqtorch.utils import ATOL, density_mat, product_state, promote_op, random_state
 
 state_000 = product_state("000")
 state_001 = product_state("001")
@@ -333,7 +333,7 @@ def test_dm(n_qubits: Tensor, batch_size: Tensor) -> None:
 
 @pytest.mark.parametrize("target,n_qubits", random_param)
 def test_promote(target: int, n_qubits: int) -> None:
-    I_prom = promote_ope(I(0).unitary(), target, n_qubits)
+    I_prom = promote_op(I(0).unitary(), target, n_qubits)
     assert I_prom.size() == torch.Size([2**n_qubits, 2**n_qubits, 1])
-    X_prom = promote_ope(X(0).unitary(), target, n_qubits)
+    X_prom = promote_op(X(0).unitary(), target, n_qubits)
     assert X_prom.size() == torch.Size([2**n_qubits, 2**n_qubits, 1])

From 62d595e8cba295f807dae07438838166984a968e Mon Sep 17 00:00:00 2001
From: EthanObadia <ethan.obadia@pasqal.com>
Date: Mon, 15 Apr 2024 18:56:02 +0200
Subject: [PATCH 10/26] Avoid circular import

---
 pyqtorch/utils.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/pyqtorch/utils.py b/pyqtorch/utils.py
index bf20fc99..2932f39a 100644
--- a/pyqtorch/utils.py
+++ b/pyqtorch/utils.py
@@ -7,7 +7,6 @@
 from torch import Tensor
 
 from pyqtorch.matrices import DEFAULT_MATRIX_DTYPE, DEFAULT_REAL_DTYPE
-from pyqtorch.primitive import I
 
 State = Tensor
 Operator = Tensor
@@ -164,12 +163,13 @@ def promote_ope(operator: Tensor, target: int, n_qubits: int) -> Tensor:
         raise ValueError("The target must be a register qubit")
 
     qubits_support = torch.arange(0, n_qubits)
+    identity = torch.tensor([[[1.0 + 0.0j], [0.0 + 0.0j]], [[0.0 + 0.0j], [1.0 + 0.0j]]])
     for support in qubits_support:
         if target > support:
-            operator = torch.kron(I(support).unitary(), operator.contiguous())
+            operator = torch.kron(identity, operator.contiguous())
         # Add .contiguous() because kron does not support the transpose (dagger)
         elif target < support:
-            operator = torch.kron(operator.contiguous(), I(support).unitary())
+            operator = torch.kron(operator.contiguous(), identity)
     return operator
 
 

From b80641064a908106d43a10ed1c6b32d628c30f0a Mon Sep 17 00:00:00 2001
From: EthanObadia <ethan.obadia@pasqal.com>
Date: Tue, 16 Apr 2024 11:04:02 +0200
Subject: [PATCH 11/26] Remove interdependency between PR

---
 pyqtorch/utils.py | 30 ------------------------------
 1 file changed, 30 deletions(-)

diff --git a/pyqtorch/utils.py b/pyqtorch/utils.py
index 2932f39a..ae9bceb9 100644
--- a/pyqtorch/utils.py
+++ b/pyqtorch/utils.py
@@ -143,36 +143,6 @@ def density_mat(state: Tensor) -> Tensor:
     undo_perm = (1, 2, 0)
     return torch.permute(torch.einsum("bi,bj->bij", (state, state.conj())), undo_perm)
 
-
-def promote_ope(operator: Tensor, target: int, n_qubits: int) -> Tensor:
-    """
-    Promotes `operator` to the size of the circuit.
-
-    Args:
-        operator (Tensor): The operator tensor to be promoted.
-        target (int): The target qubit index.
-        n_qubits (int): Number of qubits in the circuit.
-
-    Returns:
-        Tensor: The promoted operator tensor.
-
-    Raises:
-        ValueError: If `target` is not within the range of qubits.
-    """
-    if target > n_qubits:
-        raise ValueError("The target must be a register qubit")
-
-    qubits_support = torch.arange(0, n_qubits)
-    identity = torch.tensor([[[1.0 + 0.0j], [0.0 + 0.0j]], [[0.0 + 0.0j], [1.0 + 0.0j]]])
-    for support in qubits_support:
-        if target > support:
-            operator = torch.kron(identity, operator.contiguous())
-        # Add .contiguous() because kron does not support the transpose (dagger)
-        elif target < support:
-            operator = torch.kron(operator.contiguous(), identity)
-    return operator
-
-
 def batch_first(operator: Tensor) -> Tensor:
     """
     Permute the operator's batch dimension on first dimension.

From 05fa8a3a748c356a7dd7d6bc6059bb492bd7af10 Mon Sep 17 00:00:00 2001
From: EthanObadia <ethan.obadia@pasqal.com>
Date: Tue, 16 Apr 2024 14:01:42 +0200
Subject: [PATCH 12/26] Modify after review #155

---
 pyqtorch/utils.py     | 16 ++++++++--------
 tests/test_digital.py | 21 ++++++++++++---------
 2 files changed, 20 insertions(+), 17 deletions(-)

diff --git a/pyqtorch/utils.py b/pyqtorch/utils.py
index b91b44a1..ebd56900 100644
--- a/pyqtorch/utils.py
+++ b/pyqtorch/utils.py
@@ -145,6 +145,8 @@ def density_mat(state: Tensor) -> Tensor:
 
 
 def promote_op(operator: Tensor, target: int, n_qubits: int) -> Tensor:
+    from pyqtorch.primitive import I
+
     """
     Promotes `operator` to the size of the circuit (number of qubits and batch).
     Targeting the first qubit implies target = 0, so target > n_qubits - 1.
@@ -163,14 +165,12 @@ def promote_op(operator: Tensor, target: int, n_qubits: int) -> Tensor:
     """
     if target > n_qubits - 1:
         raise ValueError("The target must be a valid qubit index within the circuit's range.")
-
     qubits = torch.arange(0, n_qubits)
-    # Define I instead of importing it from pyqtorch.primitive to avoid circular import
-    identity = torch.tensor([[[1.0 + 0.0j], [0.0 + 0.0j]], [[0.0 + 0.0j], [1.0 + 0.0j]]])
+    qubits = qubits[qubits != target]
     for qubit in qubits:
-        if target > qubit:
-            operator = torch.kron(identity, operator.contiguous())
-        # Add .contiguous() because kron does not support the transpose (dagger)
-        elif target < qubit:
-            operator = torch.kron(operator.contiguous(), identity)
+        operator = torch.where(
+            target > qubit,
+            torch.kron(I(target).unitary(), operator.contiguous()),
+            torch.kron(operator.contiguous(), I(target).unitary()),
+        )
     return operator
diff --git a/tests/test_digital.py b/tests/test_digital.py
index ad5d0661..f49df716 100644
--- a/tests/test_digital.py
+++ b/tests/test_digital.py
@@ -9,10 +9,10 @@
 from torch import Tensor
 
 import pyqtorch as pyq
-from pyqtorch.apply import apply_operator
-from pyqtorch.matrices import DEFAULT_MATRIX_DTYPE, IMAT, ZMAT
+from pyqtorch.apply import apply_ope_ope, apply_operator
+from pyqtorch.matrices import DEFAULT_MATRIX_DTYPE, IMAT, ZMAT, _dagger
 from pyqtorch.parametric import Parametric
-from pyqtorch.primitive import I, X
+from pyqtorch.primitive import H, I, S, T, X, Y, Z
 from pyqtorch.utils import ATOL, density_mat, product_state, promote_op, random_state
 
 state_000 = product_state("000")
@@ -327,13 +327,16 @@ def test_dm(n_qubits: Tensor, batch_size: Tensor) -> None:
 
 size = (5, 2)
 random_param = torch.randperm(size[0] * size[1])
-random_param = random_param.view(size)
+random_param = random_param.reshape(size)
 random_param = torch.sort(random_param, dim=1)[0]
 
 
 @pytest.mark.parametrize("target,n_qubits", random_param)
-def test_promote(target: int, n_qubits: int) -> None:
-    I_prom = promote_op(I(0).unitary(), target, n_qubits)
-    assert I_prom.size() == torch.Size([2**n_qubits, 2**n_qubits, 1])
-    X_prom = promote_op(X(0).unitary(), target, n_qubits)
-    assert X_prom.size() == torch.Size([2**n_qubits, 2**n_qubits, 1])
+@pytest.mark.parametrize("operator", [I, X, Y, Z, H, T, S])
+def test_promote(target: int, n_qubits: int, operator: Tensor) -> None:
+    op_prom = promote_op(operator(target).unitary(), target, n_qubits)
+    assert op_prom.size() == torch.Size([2**n_qubits, 2**n_qubits, 1])
+    assert torch.allclose(
+        apply_ope_ope(op_prom, _dagger(op_prom), target),
+        torch.eye(2**n_qubits, dtype=torch.cdouble).unsqueeze(2),
+    )

From f0946096493ae1ea2e81efe7e3b8ade13ee4133c Mon Sep 17 00:00:00 2001
From: EthanObadia <ethan.obadia@pasqal.com>
Date: Tue, 16 Apr 2024 14:20:09 +0200
Subject: [PATCH 13/26] syntax

---
 pyqtorch/utils.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/pyqtorch/utils.py b/pyqtorch/utils.py
index 6f670408..e411e46f 100644
--- a/pyqtorch/utils.py
+++ b/pyqtorch/utils.py
@@ -146,7 +146,6 @@ def density_mat(state: Tensor) -> Tensor:
 
 def promote_op(operator: Tensor, target: int, n_qubits: int) -> Tensor:
     from pyqtorch.primitive import I
-
     """
     Promotes `operator` to the size of the circuit (number of qubits and batch).
     Targeting the first qubit implies target = 0, so target > n_qubits - 1.

From 1cd14c8e0d2c0af4401b7c838224188cb3b96c57 Mon Sep 17 00:00:00 2001
From: EthanObadia <ethan.obadia@pasqal.com>
Date: Tue, 16 Apr 2024 14:25:27 +0200
Subject: [PATCH 14/26] pre-commit

---
 pyqtorch/utils.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pyqtorch/utils.py b/pyqtorch/utils.py
index e411e46f..6f670408 100644
--- a/pyqtorch/utils.py
+++ b/pyqtorch/utils.py
@@ -146,6 +146,7 @@ def density_mat(state: Tensor) -> Tensor:
 
 def promote_op(operator: Tensor, target: int, n_qubits: int) -> Tensor:
     from pyqtorch.primitive import I
+
     """
     Promotes `operator` to the size of the circuit (number of qubits and batch).
     Targeting the first qubit implies target = 0, so target > n_qubits - 1.

From 5e5068d07703e2826418971465324be925bfe201 Mon Sep 17 00:00:00 2001
From: EthanObadia <ethan.obadia@pasqal.com>
Date: Tue, 16 Apr 2024 14:42:29 +0200
Subject: [PATCH 15/26] Rename in apply_op_op

---
 pyqtorch/apply.py     | 2 +-
 tests/test_digital.py | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/pyqtorch/apply.py b/pyqtorch/apply.py
index 8528e63c..d10ee019 100644
--- a/pyqtorch/apply.py
+++ b/pyqtorch/apply.py
@@ -61,7 +61,7 @@ def apply_operator(
     return einsum(f"{operator_dims},{in_state_dims}->{out_state_dims}", operator, state)
 
 
-def apply_ope_ope(operator_1: Tensor, operator_2: Tensor, target: int) -> Tensor:
+def apply_op_op(operator_1: Tensor, operator_2: Tensor, target: int) -> Tensor:
     """
     Compute the product of two operators.
 
diff --git a/tests/test_digital.py b/tests/test_digital.py
index f49df716..f8b0230e 100644
--- a/tests/test_digital.py
+++ b/tests/test_digital.py
@@ -9,7 +9,7 @@
 from torch import Tensor
 
 import pyqtorch as pyq
-from pyqtorch.apply import apply_ope_ope, apply_operator
+from pyqtorch.apply import apply_op_op, apply_operator
 from pyqtorch.matrices import DEFAULT_MATRIX_DTYPE, IMAT, ZMAT, _dagger
 from pyqtorch.parametric import Parametric
 from pyqtorch.primitive import H, I, S, T, X, Y, Z
@@ -337,6 +337,6 @@ def test_promote(target: int, n_qubits: int, operator: Tensor) -> None:
     op_prom = promote_op(operator(target).unitary(), target, n_qubits)
     assert op_prom.size() == torch.Size([2**n_qubits, 2**n_qubits, 1])
     assert torch.allclose(
-        apply_ope_ope(op_prom, _dagger(op_prom), target),
+        apply_op_op(op_prom, _dagger(op_prom), target),
         torch.eye(2**n_qubits, dtype=torch.cdouble).unsqueeze(2),
     )

From e1f266de50aa3ee105ca434a315cbddfa5231209 Mon Sep 17 00:00:00 2001
From: EthanObadia <ethan.obadia@pasqal.com>
Date: Tue, 16 Apr 2024 15:11:03 +0200
Subject: [PATCH 16/26] Add test_apply_op_op #153

---
 tests/test_digital.py | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/tests/test_digital.py b/tests/test_digital.py
index f8b0230e..043c4721 100644
--- a/tests/test_digital.py
+++ b/tests/test_digital.py
@@ -340,3 +340,21 @@ def test_promote(target: int, n_qubits: int, operator: Tensor) -> None:
         apply_op_op(op_prom, _dagger(op_prom), target),
         torch.eye(2**n_qubits, dtype=torch.cdouble).unsqueeze(2),
     )
+
+
+size = (3, 3)
+random_param = torch.randperm(size[0] * size[1])
+random_param = random_param.reshape(size)
+random_param = torch.sort(random_param, dim=1)[0]
+
+
+@pytest.mark.parametrize("target,n_qubits,batch_size", random_param)
+@pytest.mark.parametrize("operator", [I, X, Y, Z, H, T, S])
+def test_apply_op_op(target: int, n_qubits: int, batch_size: int, operator: Tensor) -> None:
+    op_prom: Tensor = promote_op(operator(target).unitary(), target, n_qubits)
+    op_prom = op_prom.repeat(1, 1, batch_size)
+    op_mul = apply_op_op(op_prom, _dagger(operator(target).unitary()), target)
+    assert op_mul.size() == torch.Size([2**n_qubits, 2**n_qubits, batch_size])
+    assert torch.allclose(
+        op_mul, torch.eye(2**n_qubits, dtype=torch.cdouble).unsqueeze(2).repeat(1, 1, batch_size)
+    )

From 5bea869327746507bfd99b31cd5a483db3f1b78d Mon Sep 17 00:00:00 2001
From: EthanObadia <ethan.obadia@pasqal.com>
Date: Wed, 17 Apr 2024 11:16:17 +0200
Subject: [PATCH 17/26] Modify after review

---
 pyqtorch/apply.py     | 19 ++++++++-----------
 pyqtorch/primitive.py |  6 +++---
 pyqtorch/utils.py     |  8 ++++----
 tests/test_digital.py |  5 +++--
 4 files changed, 18 insertions(+), 20 deletions(-)

diff --git a/pyqtorch/apply.py b/pyqtorch/apply.py
index d10ee019..e7d59d56 100644
--- a/pyqtorch/apply.py
+++ b/pyqtorch/apply.py
@@ -74,20 +74,17 @@ def apply_op_op(operator_1: Tensor, operator_2: Tensor, target: int) -> Tensor:
         Tensor: The product of the two operators.
     """
 
-    # Dimension verifications:
     n_qubits_1 = int(log2(operator_1.size(1)))
     n_qubits_2 = int(log2(operator_2.size(1)))
     batch_size_1 = operator_1.size(-1)
     batch_size_2 = operator_2.size(-1)
-    if n_qubits_1 != n_qubits_2:
-        if n_qubits_1 > n_qubits_2:
-            operator_2 = promote_op(operator_2, target, n_qubits_1)
-        if n_qubits_1 < n_qubits_2:
-            operator_1 = promote_op(operator_1, target, n_qubits_2)
-    if batch_size_1 != batch_size_2:
-        if batch_size_1 > batch_size_2:
-            operator_2 = operator_2.repeat(1, 1, batch_size_1)
-        if batch_size_2 > batch_size_1:
-            operator_1 = operator_1.repeat(1, 1, batch_size_2)
+    if n_qubits_1 > n_qubits_2:
+        operator_2 = promote_op(operator_2, target, n_qubits_1)
+    if n_qubits_1 < n_qubits_2:
+        operator_1 = promote_op(operator_1, target, n_qubits_2)
+    if batch_size_1 > batch_size_2:
+        operator_2 = operator_2.repeat(1, 1, batch_size_1)
+    if batch_size_2 > batch_size_1:
+        operator_1 = operator_1.repeat(1, 1, batch_size_2)
 
     return batch_last(torch.bmm(batch_first(operator_1), batch_first(operator_2)))
diff --git a/pyqtorch/primitive.py b/pyqtorch/primitive.py
index f6683377..786100b9 100644
--- a/pyqtorch/primitive.py
+++ b/pyqtorch/primitive.py
@@ -41,15 +41,15 @@ def extra_repr(self) -> str:
     def param_type(self) -> None:
         return self._param_type
 
-    def unitary(self, values: dict[str, Tensor] | Tensor = {}) -> Tensor:
+    def unitary(self, values: dict[str, Tensor] | Tensor = dict()) -> Tensor:
         return self.pauli.unsqueeze(2)
 
-    def forward(self, state: Tensor, values: dict[str, Tensor] | Tensor = {}) -> Tensor:
+    def forward(self, state: Tensor, values: dict[str, Tensor] | Tensor = dict()) -> Tensor:
         return apply_operator(
             state, self.unitary(values), self.qubit_support, len(state.size()) - 1
         )
 
-    def dagger(self, values: dict[str, Tensor] | Tensor = {}) -> Tensor:
+    def dagger(self, values: dict[str, Tensor] | Tensor = dict()) -> Tensor:
         return _dagger(self.unitary(values))
 
     @property
diff --git a/pyqtorch/utils.py b/pyqtorch/utils.py
index 6f670408..4120bd03 100644
--- a/pyqtorch/utils.py
+++ b/pyqtorch/utils.py
@@ -181,10 +181,10 @@ def batch_first(operator: Tensor) -> Tensor:
     Permute the operator's batch dimension on first dimension.
 
     Args:
-    operator (Tensor): Operator in size [2**n_qubits, 2**n_qubits,batch_size].
+        operator (Tensor): Operator in size [2**n_qubits, 2**n_qubits,batch_size].
 
     Returns:
-    Tensor: Operator in size [batch_size, 2**n_qubits, 2**n_qubits].
+        Tensor: Operator in size [batch_size, 2**n_qubits, 2**n_qubits].
     """
     batch_first_perm = (2, 0, 1)
     return torch.permute(operator, batch_first_perm)
@@ -195,10 +195,10 @@ def batch_last(operator: Tensor) -> Tensor:
     Permute the operator's batch dimension on last dimension.
 
     Args:
-    operator (Tensor): Operator in size [batch_size,2**n_qubits, 2**n_qubits].
+        operator (Tensor): Operator in size [batch_size,2**n_qubits, 2**n_qubits].
 
     Returns:
-    Tensor: Operator in size [2**n_qubits, 2**n_qubits,batch_size].
+        Tensor: Operator in size [2**n_qubits, 2**n_qubits,batch_size].
     """
     undo_perm = (1, 2, 0)
     return torch.permute(operator, undo_perm)
diff --git a/tests/test_digital.py b/tests/test_digital.py
index 043c4721..b06a8824 100644
--- a/tests/test_digital.py
+++ b/tests/test_digital.py
@@ -329,10 +329,11 @@ def test_dm(n_qubits: Tensor, batch_size: Tensor) -> None:
 random_param = torch.randperm(size[0] * size[1])
 random_param = random_param.reshape(size)
 random_param = torch.sort(random_param, dim=1)[0]
+GATESET = [I, X, Y, Z, H, T, S]
 
 
 @pytest.mark.parametrize("target,n_qubits", random_param)
-@pytest.mark.parametrize("operator", [I, X, Y, Z, H, T, S])
+@pytest.mark.parametrize("operator", GATESET)
 def test_promote(target: int, n_qubits: int, operator: Tensor) -> None:
     op_prom = promote_op(operator(target).unitary(), target, n_qubits)
     assert op_prom.size() == torch.Size([2**n_qubits, 2**n_qubits, 1])
@@ -349,7 +350,7 @@ def test_promote(target: int, n_qubits: int, operator: Tensor) -> None:
 
 
 @pytest.mark.parametrize("target,n_qubits,batch_size", random_param)
-@pytest.mark.parametrize("operator", [I, X, Y, Z, H, T, S])
+@pytest.mark.parametrize("operator", GATESET)
 def test_apply_op_op(target: int, n_qubits: int, batch_size: int, operator: Tensor) -> None:
     op_prom: Tensor = promote_op(operator(target).unitary(), target, n_qubits)
     op_prom = op_prom.repeat(1, 1, batch_size)

From b1d15e64c062efa66715a06aea9965532ae064f2 Mon Sep 17 00:00:00 2001
From: EthanObadia <ethan.obadia@pasqal.com>
Date: Wed, 17 Apr 2024 20:37:23 +0200
Subject: [PATCH 18/26] Modify after review #153

---
 pyqtorch/apply.py     | 12 ++++++------
 pyqtorch/utils.py     |  2 +-
 tests/test_digital.py | 38 +++++++++++++++-----------------------
 3 files changed, 22 insertions(+), 30 deletions(-)

diff --git a/pyqtorch/apply.py b/pyqtorch/apply.py
index e7d59d56..b2005899 100644
--- a/pyqtorch/apply.py
+++ b/pyqtorch/apply.py
@@ -9,7 +9,7 @@
 from numpy.typing import NDArray
 from torch import Tensor, einsum
 
-from pyqtorch.utils import batch_first, batch_last, promote_op
+from pyqtorch.utils import batch_first, batch_last, promote_operator
 
 ABC_ARRAY: NDArray = array(list(ABC))
 
@@ -61,7 +61,7 @@ def apply_operator(
     return einsum(f"{operator_dims},{in_state_dims}->{out_state_dims}", operator, state)
 
 
-def apply_op_op(operator_1: Tensor, operator_2: Tensor, target: int) -> Tensor:
+def operator_product(operator_1: Tensor, operator_2: Tensor, target: int) -> Tensor:
     """
     Compute the product of two operators.
 
@@ -79,12 +79,12 @@ def apply_op_op(operator_1: Tensor, operator_2: Tensor, target: int) -> Tensor:
     batch_size_1 = operator_1.size(-1)
     batch_size_2 = operator_2.size(-1)
     if n_qubits_1 > n_qubits_2:
-        operator_2 = promote_op(operator_2, target, n_qubits_1)
-    if n_qubits_1 < n_qubits_2:
-        operator_1 = promote_op(operator_1, target, n_qubits_2)
+        operator_2 = promote_operator(operator_2, target, n_qubits_1)
+    elif n_qubits_1 < n_qubits_2:
+        operator_1 = promote_operator(operator_1, target, n_qubits_2)
     if batch_size_1 > batch_size_2:
         operator_2 = operator_2.repeat(1, 1, batch_size_1)
-    if batch_size_2 > batch_size_1:
+    elif batch_size_2 > batch_size_1:
         operator_1 = operator_1.repeat(1, 1, batch_size_2)
 
     return batch_last(torch.bmm(batch_first(operator_1), batch_first(operator_2)))
diff --git a/pyqtorch/utils.py b/pyqtorch/utils.py
index 4120bd03..a3ef3bd6 100644
--- a/pyqtorch/utils.py
+++ b/pyqtorch/utils.py
@@ -144,7 +144,7 @@ def density_mat(state: Tensor) -> Tensor:
     return torch.permute(torch.einsum("bi,bj->bij", (state, state.conj())), undo_perm)
 
 
-def promote_op(operator: Tensor, target: int, n_qubits: int) -> Tensor:
+def promote_operator(operator: Tensor, target: int, n_qubits: int) -> Tensor:
     from pyqtorch.primitive import I
 
     """
diff --git a/tests/test_digital.py b/tests/test_digital.py
index b06a8824..eee8c6f0 100644
--- a/tests/test_digital.py
+++ b/tests/test_digital.py
@@ -9,11 +9,11 @@
 from torch import Tensor
 
 import pyqtorch as pyq
-from pyqtorch.apply import apply_op_op, apply_operator
+from pyqtorch.apply import apply_operator, operator_product
 from pyqtorch.matrices import DEFAULT_MATRIX_DTYPE, IMAT, ZMAT, _dagger
 from pyqtorch.parametric import Parametric
 from pyqtorch.primitive import H, I, S, T, X, Y, Z
-from pyqtorch.utils import ATOL, density_mat, product_state, promote_op, random_state
+from pyqtorch.utils import ATOL, density_mat, product_state, promote_operator, random_state
 
 state_000 = product_state("000")
 state_001 = product_state("001")
@@ -25,6 +25,8 @@
 state_1110 = product_state("1110")
 state_1111 = product_state("1111")
 
+GATESET = [I, X, Y, Z, H, T, S]
+
 
 def test_identity() -> None:
     assert torch.allclose(product_state("0"), pyq.I(0)(product_state("0"), None))
@@ -325,36 +327,26 @@ def test_dm(n_qubits: Tensor, batch_size: Tensor) -> None:
     assert torch.allclose(dm, dm_proj)
 
 
-size = (5, 2)
-random_param = torch.randperm(size[0] * size[1])
-random_param = random_param.reshape(size)
-random_param = torch.sort(random_param, dim=1)[0]
-GATESET = [I, X, Y, Z, H, T, S]
-
-
-@pytest.mark.parametrize("target,n_qubits", random_param)
 @pytest.mark.parametrize("operator", GATESET)
-def test_promote(target: int, n_qubits: int, operator: Tensor) -> None:
-    op_prom = promote_op(operator(target).unitary(), target, n_qubits)
+def test_promote(operator: Tensor) -> None:
+    n_qubits = torch.randint(low=1, high=8, size=(1,)).item()
+    target = random.choice([i for i in range(n_qubits)])
+    op_prom = promote_operator(operator(target).unitary(), target, n_qubits)
     assert op_prom.size() == torch.Size([2**n_qubits, 2**n_qubits, 1])
     assert torch.allclose(
-        apply_op_op(op_prom, _dagger(op_prom), target),
+        operator_product(op_prom, _dagger(op_prom), target),
         torch.eye(2**n_qubits, dtype=torch.cdouble).unsqueeze(2),
     )
 
 
-size = (3, 3)
-random_param = torch.randperm(size[0] * size[1])
-random_param = random_param.reshape(size)
-random_param = torch.sort(random_param, dim=1)[0]
-
-
-@pytest.mark.parametrize("target,n_qubits,batch_size", random_param)
 @pytest.mark.parametrize("operator", GATESET)
-def test_apply_op_op(target: int, n_qubits: int, batch_size: int, operator: Tensor) -> None:
-    op_prom: Tensor = promote_op(operator(target).unitary(), target, n_qubits)
+def test_operator_product(operator: Tensor) -> None:
+    n_qubits = torch.randint(low=1, high=8, size=(1,)).item()
+    target = random.choice([i for i in range(n_qubits)])
+    batch_size = torch.randint(low=1, high=5, size=(1,)).item()
+    op_prom: Tensor = promote_operator(operator(target).unitary(), target, n_qubits)
     op_prom = op_prom.repeat(1, 1, batch_size)
-    op_mul = apply_op_op(op_prom, _dagger(operator(target).unitary()), target)
+    op_mul = operator_product(operator(target).unitary(), _dagger(op_prom), target)
     assert op_mul.size() == torch.Size([2**n_qubits, 2**n_qubits, batch_size])
     assert torch.allclose(
         op_mul, torch.eye(2**n_qubits, dtype=torch.cdouble).unsqueeze(2).repeat(1, 1, batch_size)

From f10f29fa2896e95db4e6233ca55a2aaf612a9676 Mon Sep 17 00:00:00 2001
From: EthanObadia <ethan.obadia@pasqal.com>
Date: Fri, 19 Apr 2024 17:04:00 +0200
Subject: [PATCH 19/26] Modify operator_product and test #153

---
 pyqtorch/apply.py     | 24 ++++++++++++------------
 tests/test_digital.py | 17 +++++++++++------
 2 files changed, 23 insertions(+), 18 deletions(-)

diff --git a/pyqtorch/apply.py b/pyqtorch/apply.py
index b2005899..746b8af9 100644
--- a/pyqtorch/apply.py
+++ b/pyqtorch/apply.py
@@ -61,30 +61,30 @@ def apply_operator(
     return einsum(f"{operator_dims},{in_state_dims}->{out_state_dims}", operator, state)
 
 
-def operator_product(operator_1: Tensor, operator_2: Tensor, target: int) -> Tensor:
+def operator_product(op1: Tensor, op2: Tensor, target: int) -> Tensor:
     """
     Compute the product of two operators.
 
     Args:
-        operator_1 (Tensor): The first operator.
-        operator_2 (Tensor): The second operator.
+        op1 (Tensor): The first operator.
+        op2 (Tensor): The second operator.
         target (int): The target qubit index.
 
     Returns:
         Tensor: The product of the two operators.
     """
 
-    n_qubits_1 = int(log2(operator_1.size(1)))
-    n_qubits_2 = int(log2(operator_2.size(1)))
-    batch_size_1 = operator_1.size(-1)
-    batch_size_2 = operator_2.size(-1)
+    n_qubits_1 = int(log2(op1.size(1)))
+    n_qubits_2 = int(log2(op2.size(1)))
+    batch_size_1 = op1.size(-1)
+    batch_size_2 = op2.size(-1)
     if n_qubits_1 > n_qubits_2:
-        operator_2 = promote_operator(operator_2, target, n_qubits_1)
+        op2 = promote_operator(op2, target, n_qubits_1)
     elif n_qubits_1 < n_qubits_2:
-        operator_1 = promote_operator(operator_1, target, n_qubits_2)
+        op1 = promote_operator(op1, target, n_qubits_2)
     if batch_size_1 > batch_size_2:
-        operator_2 = operator_2.repeat(1, 1, batch_size_1)
+        op2 = op2.repeat(1, 1, batch_size_1)[:, :, :batch_size_1]
     elif batch_size_2 > batch_size_1:
-        operator_1 = operator_1.repeat(1, 1, batch_size_2)
+        op1 = op1.repeat(1, 1, batch_size_2)[:, :, :batch_size_2]
 
-    return batch_last(torch.bmm(batch_first(operator_1), batch_first(operator_2)))
+    return batch_last(torch.bmm(batch_first(op1), batch_first(op2)))
diff --git a/tests/test_digital.py b/tests/test_digital.py
index eee8c6f0..7f09e85f 100644
--- a/tests/test_digital.py
+++ b/tests/test_digital.py
@@ -343,11 +343,16 @@ def test_promote(operator: Tensor) -> None:
 def test_operator_product(operator: Tensor) -> None:
     n_qubits = torch.randint(low=1, high=8, size=(1,)).item()
     target = random.choice([i for i in range(n_qubits)])
-    batch_size = torch.randint(low=1, high=5, size=(1,)).item()
-    op_prom: Tensor = promote_operator(operator(target).unitary(), target, n_qubits)
-    op_prom = op_prom.repeat(1, 1, batch_size)
-    op_mul = operator_product(operator(target).unitary(), _dagger(op_prom), target)
-    assert op_mul.size() == torch.Size([2**n_qubits, 2**n_qubits, batch_size])
+    batch_size_1 = torch.randint(low=1, high=5, size=(1,)).item()
+    batch_size_2 = torch.randint(low=1, high=5, size=(1,)).item()
+    max_batch = max(batch_size_2, batch_size_1)
+    op_prom: Tensor = promote_operator(operator(target).unitary(), target, n_qubits).repeat(
+        1, 1, batch_size_1
+    )
+    op_mul = operator_product(
+        operator(target).unitary().repeat(1, 1, batch_size_2), _dagger(op_prom), target
+    )
+    assert op_mul.size() == torch.Size([2**n_qubits, 2**n_qubits, max_batch])
     assert torch.allclose(
-        op_mul, torch.eye(2**n_qubits, dtype=torch.cdouble).unsqueeze(2).repeat(1, 1, batch_size)
+        op_mul, torch.eye(2**n_qubits, dtype=torch.cdouble).unsqueeze(2).repeat(1, 1, max_batch)
     )

From 3c6d51b0964f71de84e80379f746397ca8dc08fd Mon Sep 17 00:00:00 2001
From: EthanObadia <ethan.obadia@pasqal.com>
Date: Fri, 19 Apr 2024 17:59:00 +0200
Subject: [PATCH 20/26] Modify type annotations

---
 tests/test_digital.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/test_digital.py b/tests/test_digital.py
index 7f09e85f..3a9edd08 100644
--- a/tests/test_digital.py
+++ b/tests/test_digital.py
@@ -2,7 +2,7 @@
 
 import random
 from math import log2
-from typing import Callable, Tuple
+from typing import Callable, Tuple, List
 
 import pytest
 import torch
@@ -328,7 +328,7 @@ def test_dm(n_qubits: Tensor, batch_size: Tensor) -> None:
 
 
 @pytest.mark.parametrize("operator", GATESET)
-def test_promote(operator: Tensor) -> None:
+def test_promote(operator: List) -> None:
     n_qubits = torch.randint(low=1, high=8, size=(1,)).item()
     target = random.choice([i for i in range(n_qubits)])
     op_prom = promote_operator(operator(target).unitary(), target, n_qubits)
@@ -340,7 +340,7 @@ def test_promote(operator: Tensor) -> None:
 
 
 @pytest.mark.parametrize("operator", GATESET)
-def test_operator_product(operator: Tensor) -> None:
+def test_operator_product(operator: List) -> None:
     n_qubits = torch.randint(low=1, high=8, size=(1,)).item()
     target = random.choice([i for i in range(n_qubits)])
     batch_size_1 = torch.randint(low=1, high=5, size=(1,)).item()

From 7442e1255a51c829980828f35b20062447d75037 Mon Sep 17 00:00:00 2001
From: EthanObadia <ethan.obadia@pasqal.com>
Date: Fri, 19 Apr 2024 18:23:14 +0200
Subject: [PATCH 21/26] Modify type annotation

---
 tests/test_digital.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/test_digital.py b/tests/test_digital.py
index 3a9edd08..107b28ec 100644
--- a/tests/test_digital.py
+++ b/tests/test_digital.py
@@ -2,7 +2,7 @@
 
 import random
 from math import log2
-from typing import Callable, Tuple, List
+from typing import Callable, Tuple
 
 import pytest
 import torch
@@ -328,7 +328,7 @@ def test_dm(n_qubits: Tensor, batch_size: Tensor) -> None:
 
 
 @pytest.mark.parametrize("operator", GATESET)
-def test_promote(operator: List) -> None:
+def test_promote(operator: Parametric) -> None:
     n_qubits = torch.randint(low=1, high=8, size=(1,)).item()
     target = random.choice([i for i in range(n_qubits)])
     op_prom = promote_operator(operator(target).unitary(), target, n_qubits)
@@ -340,7 +340,7 @@ def test_promote(operator: List) -> None:
 
 
 @pytest.mark.parametrize("operator", GATESET)
-def test_operator_product(operator: List) -> None:
+def test_operator_product(operator: Parametric) -> None:
     n_qubits = torch.randint(low=1, high=8, size=(1,)).item()
     target = random.choice([i for i in range(n_qubits)])
     batch_size_1 = torch.randint(low=1, high=5, size=(1,)).item()

From 23ce1c63add8e319b5bb2f75db16dcaeef8acd38 Mon Sep 17 00:00:00 2001
From: EthanObadia <ethan.obadia@pasqal.com>
Date: Mon, 22 Apr 2024 11:40:09 +0200
Subject: [PATCH 22/26] Use fixture for test #153

---
 tests/test_digital.py | 25 ++++++++++++-------------
 1 file changed, 12 insertions(+), 13 deletions(-)

diff --git a/tests/test_digital.py b/tests/test_digital.py
index 107b28ec..0c03ffb6 100644
--- a/tests/test_digital.py
+++ b/tests/test_digital.py
@@ -2,7 +2,7 @@
 
 import random
 from math import log2
-from typing import Callable, Tuple
+from typing import Any, Callable, Tuple
 
 import pytest
 import torch
@@ -12,7 +12,7 @@
 from pyqtorch.apply import apply_operator, operator_product
 from pyqtorch.matrices import DEFAULT_MATRIX_DTYPE, IMAT, ZMAT, _dagger
 from pyqtorch.parametric import Parametric
-from pyqtorch.primitive import H, I, S, T, X, Y, Z
+from pyqtorch.primitive import H, I, Primitive, S, T, X, Y, Z
 from pyqtorch.utils import ATOL, density_mat, product_state, promote_operator, random_state
 
 state_000 = product_state("000")
@@ -25,8 +25,6 @@
 state_1110 = product_state("1110")
 state_1111 = product_state("1111")
 
-GATESET = [I, X, Y, Z, H, T, S]
-
 
 def test_identity() -> None:
     assert torch.allclose(product_state("0"), pyq.I(0)(product_state("0"), None))
@@ -327,11 +325,15 @@ def test_dm(n_qubits: Tensor, batch_size: Tensor) -> None:
     assert torch.allclose(dm, dm_proj)
 
 
-@pytest.mark.parametrize("operator", GATESET)
-def test_promote(operator: Parametric) -> None:
+@pytest.fixture(params=[I, X, Y, Z, H, T, S])
+def GATE(request: Primitive) -> Any:
+    return request.param
+
+
+def test_promote(GATE: Primitive) -> None:
     n_qubits = torch.randint(low=1, high=8, size=(1,)).item()
     target = random.choice([i for i in range(n_qubits)])
-    op_prom = promote_operator(operator(target).unitary(), target, n_qubits)
+    op_prom = promote_operator(GATE(target).unitary(), target, n_qubits)
     assert op_prom.size() == torch.Size([2**n_qubits, 2**n_qubits, 1])
     assert torch.allclose(
         operator_product(op_prom, _dagger(op_prom), target),
@@ -339,18 +341,15 @@ def test_promote(operator: Parametric) -> None:
     )
 
 
-@pytest.mark.parametrize("operator", GATESET)
-def test_operator_product(operator: Parametric) -> None:
+def test_operator_product(GATE: Primitive) -> None:
     n_qubits = torch.randint(low=1, high=8, size=(1,)).item()
     target = random.choice([i for i in range(n_qubits)])
     batch_size_1 = torch.randint(low=1, high=5, size=(1,)).item()
     batch_size_2 = torch.randint(low=1, high=5, size=(1,)).item()
     max_batch = max(batch_size_2, batch_size_1)
-    op_prom: Tensor = promote_operator(operator(target).unitary(), target, n_qubits).repeat(
-        1, 1, batch_size_1
-    )
+    op_prom = promote_operator(GATE(target).unitary(), target, n_qubits).repeat(1, 1, batch_size_1)
     op_mul = operator_product(
-        operator(target).unitary().repeat(1, 1, batch_size_2), _dagger(op_prom), target
+        GATE(target).unitary().repeat(1, 1, batch_size_2), _dagger(op_prom), target
     )
     assert op_mul.size() == torch.Size([2**n_qubits, 2**n_qubits, max_batch])
     assert torch.allclose(

From e4bf0efafea24f89fe106a5ba1e4b8f845294c93 Mon Sep 17 00:00:00 2001
From: EthanObadia <ethan.obadia@pasqal.com>
Date: Mon, 22 Apr 2024 14:32:03 +0200
Subject: [PATCH 23/26] Use numpy log2

---
 pyqtorch/apply.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/pyqtorch/apply.py b/pyqtorch/apply.py
index 746b8af9..9589f4fe 100644
--- a/pyqtorch/apply.py
+++ b/pyqtorch/apply.py
@@ -1,11 +1,10 @@
 from __future__ import annotations
 
-from math import log2
 from string import ascii_letters as ABC
 from typing import Tuple
 
 import torch
-from numpy import array
+from numpy import array, log2
 from numpy.typing import NDArray
 from torch import Tensor, einsum
 

From dacf70c4b82152923abf557e82eaf3a1052d44e1 Mon Sep 17 00:00:00 2001
From: EthanObadia <ethan.obadia@pasqal.com>
Date: Mon, 22 Apr 2024 17:31:32 +0200
Subject: [PATCH 24/26] Remove typing Tuple

---
 pyqtorch/apply.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/pyqtorch/apply.py b/pyqtorch/apply.py
index 9589f4fe..2f10a2ad 100644
--- a/pyqtorch/apply.py
+++ b/pyqtorch/apply.py
@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 from string import ascii_letters as ABC
-from typing import Tuple
 
 import torch
 from numpy import array, log2
@@ -16,7 +15,7 @@
 def apply_operator(
     state: Tensor,
     operator: Tensor,
-    qubits: Tuple[int, ...] | list[int],
+    qubits: tuple[int, ...] | list[int],
     n_qubits: int = None,
     batch_size: int = None,
 ) -> Tensor:

From e96e7ec10d2185d08c4a53b5251a6d2fc0375c03 Mon Sep 17 00:00:00 2001
From: EthanObadia <ethan.obadia@pasqal.com>
Date: Wed, 24 Apr 2024 15:29:23 +0200
Subject: [PATCH 25/26] Modify test after review

---
 pyqtorch/primitive.py | 16 ++++++++--------
 tests/conftest.py     |  8 ++++++++
 tests/test_digital.py | 27 +++++++--------------------
 3 files changed, 23 insertions(+), 28 deletions(-)
 create mode 100644 tests/conftest.py

diff --git a/pyqtorch/primitive.py b/pyqtorch/primitive.py
index 786100b9..af8b9b70 100644
--- a/pyqtorch/primitive.py
+++ b/pyqtorch/primitive.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from math import log2
-from typing import Any, Tuple
+from typing import Any
 
 import torch
 from torch import Tensor
@@ -15,7 +15,7 @@ class Primitive(torch.nn.Module):
     def __init__(self, pauli: Tensor, target: int) -> None:
         super().__init__()
         self.target: int = target
-        self.qubit_support: Tuple[int, ...] = (target,)
+        self.qubit_support: tuple[int, ...] = (target,)
         self.n_qubits: int = max(self.qubit_support)
         self.register_buffer("pauli", pauli)
         self._param_type = None
@@ -136,7 +136,7 @@ def __init__(self, control: int, target: int):
 
 
 class CSWAP(Primitive):
-    def __init__(self, control: int | Tuple[int, ...], target: int):
+    def __init__(self, control: int | tuple[int, ...], target: int):
         super().__init__(OPERATIONS_DICT["CSWAP"], target)
         self.control = (control,) if isinstance(control, int) else control
         self.target = target
@@ -145,7 +145,7 @@ def __init__(self, control: int | Tuple[int, ...], target: int):
 
 
 class ControlledOperationGate(Primitive):
-    def __init__(self, gate: str, control: int | Tuple[int, ...], target: int):
+    def __init__(self, gate: str, control: int | tuple[int, ...], target: int):
         self.control = (control,) if isinstance(control, int) else control
         mat = OPERATIONS_DICT[gate]
         mat = _controlled(
@@ -159,7 +159,7 @@ def __init__(self, gate: str, control: int | Tuple[int, ...], target: int):
 
 
 class CNOT(ControlledOperationGate):
-    def __init__(self, control: int | Tuple[int, ...], target: int):
+    def __init__(self, control: int | tuple[int, ...], target: int):
         super().__init__("X", control, target)
 
 
@@ -167,15 +167,15 @@ def __init__(self, control: int | Tuple[int, ...], target: int):
 
 
 class CY(ControlledOperationGate):
-    def __init__(self, control: int | Tuple[int, ...], target: int):
+    def __init__(self, control: int | tuple[int, ...], target: int):
         super().__init__("Y", control, target)
 
 
 class CZ(ControlledOperationGate):
-    def __init__(self, control: int | Tuple[int, ...], target: int):
+    def __init__(self, control: int | tuple[int, ...], target: int):
         super().__init__("Z", control, target)
 
 
 class Toffoli(ControlledOperationGate):
-    def __init__(self, control: int | Tuple[int, ...], target: int):
+    def __init__(self, control: int | tuple[int, ...], target: int):
         super().__init__("X", control, target)
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 00000000..05fa6ff8
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,8 @@
+import pytest
+from typing import Any 
+
+from pyqtorch.primitive import I, X, Y, Z, H, T, S, Primitive
+
+@pytest.fixture(params=[I, X, Y, Z, H, T, S])
+def gate(request: Primitive) -> Any:
+    return request.param
\ No newline at end of file
diff --git a/tests/test_digital.py b/tests/test_digital.py
index 0c03ffb6..496f721a 100644
--- a/tests/test_digital.py
+++ b/tests/test_digital.py
@@ -2,7 +2,7 @@
 
 import random
 from math import log2
-from typing import Any, Callable, Tuple
+from typing import Callable, Tuple
 
 import pytest
 import torch
@@ -12,7 +12,7 @@
 from pyqtorch.apply import apply_operator, operator_product
 from pyqtorch.matrices import DEFAULT_MATRIX_DTYPE, IMAT, ZMAT, _dagger
 from pyqtorch.parametric import Parametric
-from pyqtorch.primitive import H, I, Primitive, S, T, X, Y, Z
+from pyqtorch.primitive import Primitive
 from pyqtorch.utils import ATOL, density_mat, product_state, promote_operator, random_state
 
 state_000 = product_state("000")
@@ -294,7 +294,6 @@ def test_U() -> None:
 
 @pytest.mark.parametrize("n_qubits,batch_size", torch.randint(1, 6, (8, 2)))
 def test_dm(n_qubits: Tensor, batch_size: Tensor) -> None:
-    # Test without batches:
     state = random_state(n_qubits)
     projector = torch.outer(state.flatten(), state.conj().flatten()).view(
         2**n_qubits, 2**n_qubits, 1
@@ -302,38 +301,26 @@ def test_dm(n_qubits: Tensor, batch_size: Tensor) -> None:
     dm = density_mat(state)
     assert dm.size() == torch.Size([2**n_qubits, 2**n_qubits, 1])
     assert torch.allclose(dm, projector)
-
-    # Test with batches:
     states = []
     projectors = []
-    # Batches creation:
     for batch in range(batch_size):
-        # Batch state creation:
         state = random_state(n_qubits)
         states.append(state)
-        # Batch projector:
         projector = torch.outer(state.flatten(), state.conj().flatten()).view(
             2**n_qubits, 2**n_qubits, 1
         )
         projectors.append(projector)
-    # Concatenate all the batch projectors:
     dm_proj = torch.cat(projectors, dim=2)
-    # Concatenate the batch state to compute the density matrix
     state_cat = torch.cat(states, dim=n_qubits)
     dm = density_mat(state_cat)
     assert dm.size() == torch.Size([2**n_qubits, 2**n_qubits, batch_size])
     assert torch.allclose(dm, dm_proj)
 
 
-@pytest.fixture(params=[I, X, Y, Z, H, T, S])
-def GATE(request: Primitive) -> Any:
-    return request.param
-
-
-def test_promote(GATE: Primitive) -> None:
+def test_promote(gate: Primitive) -> None:
     n_qubits = torch.randint(low=1, high=8, size=(1,)).item()
     target = random.choice([i for i in range(n_qubits)])
-    op_prom = promote_operator(GATE(target).unitary(), target, n_qubits)
+    op_prom = promote_operator(gate(target).unitary(), target, n_qubits)
     assert op_prom.size() == torch.Size([2**n_qubits, 2**n_qubits, 1])
     assert torch.allclose(
         operator_product(op_prom, _dagger(op_prom), target),
@@ -341,15 +328,15 @@ def test_promote(GATE: Primitive) -> None:
     )
 
 
-def test_operator_product(GATE: Primitive) -> None:
+def test_operator_product(gate: Primitive) -> None:
     n_qubits = torch.randint(low=1, high=8, size=(1,)).item()
     target = random.choice([i for i in range(n_qubits)])
     batch_size_1 = torch.randint(low=1, high=5, size=(1,)).item()
     batch_size_2 = torch.randint(low=1, high=5, size=(1,)).item()
     max_batch = max(batch_size_2, batch_size_1)
-    op_prom = promote_operator(GATE(target).unitary(), target, n_qubits).repeat(1, 1, batch_size_1)
+    op_prom = promote_operator(gate(target).unitary(), target, n_qubits).repeat(1, 1, batch_size_1)
     op_mul = operator_product(
-        GATE(target).unitary().repeat(1, 1, batch_size_2), _dagger(op_prom), target
+        gate(target).unitary().repeat(1, 1, batch_size_2), _dagger(op_prom), target
     )
     assert op_mul.size() == torch.Size([2**n_qubits, 2**n_qubits, max_batch])
     assert torch.allclose(

From be7c562f7815db2c656e47ce285c58f676af4112 Mon Sep 17 00:00:00 2001
From: EthanObadia <ethan.obadia@pasqal.com>
Date: Wed, 24 Apr 2024 15:34:47 +0200
Subject: [PATCH 26/26] pre commit

---
 tests/conftest.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index 05fa6ff8..632824d9 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,8 +1,12 @@
+from __future__ import annotations
+
+from typing import Any
+
 import pytest
-from typing import Any 
 
-from pyqtorch.primitive import I, X, Y, Z, H, T, S, Primitive
+from pyqtorch.primitive import H, I, Primitive, S, T, X, Y, Z
+
 
 @pytest.fixture(params=[I, X, Y, Z, H, T, S])
 def gate(request: Primitive) -> Any:
-    return request.param
\ No newline at end of file
+    return request.param