feat: added CustomCupyMatrices

src/qibojit/backends/gpu.py

@@ -5,7 +5,11 @@
 from qibo.config import log, raise_error
 
 from qibojit.backends.cpu import NumbaBackend
-from qibojit.backends.matrices import CupyMatrices, CuQuantumMatrices, CustomMatrices
+from qibojit.backends.matrices import (
+    CupyMatrices,
+    CuQuantumMatrices,
+    CustomCupyMatrices,
+)
 
 
 class CupyBackend(NumbaBackend):  # pragma: no cover
@@ -44,7 +48,7 @@ def __init__(self):
         self.device = "/GPU:0"
         self.kernel_type = "double"
         self.matrices = CupyMatrices(self.dtype)
-        self.custom_matrices = CustomMatrices(self.dtype)
+        self.custom_matrices = CustomCupyMatrices(self.dtype)
         try:
             if not cp.cuda.runtime.getDeviceCount():  # pragma: no cover
                 raise RuntimeError("Cannot use cupy backend if GPU is not available.")

src/qibojit/backends/matrices.py

@@ -23,7 +23,7 @@ def _cast(self, x, dtype):
     def Unitary(self, u):
         if isinstance(u, self.cp.ndarray):
             u = u.get()
-        return super().Unitary(u)
+        return self.super().Unitary(u)
 
 
 class CuQuantumMatrices(NumpyMatrices):
@@ -97,3 +97,35 @@ def fSim(self, theta, phi):
     def GeneralizedfSim(self, u, phi):
         phase = np.exp(-1j * phi)
         return np.array([u[0, 0], u[0, 1], u[1, 0], u[1, 1], phase], dtype=self.dtype)
+
+
+class CustomCupyMatrices(CustomMatrices):
+
+    def __init__(self, dtype):
+        super().__init__(dtype)
+        import cupy as cp
+
+        self.cp = cp
+
+    def _cast(self, x, dtype):
+        if any(isinstance(item, self.cp.ndarray) for sublist in x for item in sublist):
+            return self.cp.hstack(
+                self.cp.array(item, dtype) for sublist in x for item in sublist
+            ).reshape(len(x), len(x))
+        return self.cp.array(x, dtype=dtype)
+
+    def U1(self, theta):
+        dtype = getattr(self.cp, self.dtype)
+        return dtype(self.cp.exp(1j * theta))
+
+    def fSim(self, theta, phi):
+        cost = self.cp.cos(theta) + 0j
+        isint = -1j * self.cp.sin(theta)
+        phase = self.cp.exp(-1j * phi)
+        return self.cp.array([cost, isint, isint, cost, phase], dtype=self.dtype)
+
+    def GeneralizedfSim(self, u, phi):
+        phase = self.cp.exp(-1j * phi)
+        return self.cp.array(
+            [u[0, 0], u[0, 1], u[1, 0], u[1, 1], phase], dtype=self.dtype
+        )