feat: overwriting self.np.zeros + fix cast in PSR

Author: BrunoLiegiBastonLiegi

src/qiboml/backends/pytorch.py

@@ -71,6 +71,18 @@ def __init__(self):
         self.np.sign = self.np.sgn
         self.np.flatnonzero = lambda x: self.np.nonzero(x).flatten()
 
+        # These functions are device dependent
+        torch_zeros = self.np.zeros
+
+        def zeros(shape, dtype=None, device=None):
+            if dtype is None:
+                dtype = self.dtype
+            if device is None:
+                device = self.device
+            return torch_zeros(shape, dtype=dtype, device=device)
+
+        setattr(self.np, "zeros", zeros)
+
     def _torch_dtype(self, dtype):
         if dtype == "float":
             dtype += "32"
@@ -178,16 +190,6 @@ def _cast_parameter(self, x, trainable):
             x, dtype=self.dtype, requires_grad=trainable, device=self.device
         )
 
-    def zero_state(self, nqubits):
-        state = self.np.zeros(2**nqubits, dtype=self.dtype, device=self.device)
-        state[0] = 1
-        return state
-
-    def zero_density_matrix(self, nqubits):
-        state = self.np.zeros(2 * (2**nqubits,), dtype=self.dtype, device=self.device)
-        state[0, 0] = 1
-        return state
-
     def is_sparse(self, x):
         if isinstance(x, self.np.Tensor):
             return x.is_sparse

src/qiboml/interfaces/pytorch.py

@@ -150,14 +150,17 @@ def forward(
         ctx.backend = backend
         ctx.differentiation = differentiation
         x_clone = x.clone().detach().cpu().numpy()
-        x_clone = backend.cast(x_clone)
+        x_clone = backend.cast(x_clone, dtype=x_clone.dtype)
         params = [
-            backend.cast(par.clone().detach().cpu().numpy()) for par in parameters
+            backend.cast(par.clone().detach().cpu().numpy(), dtype=x_clone.dtype)
+            for par in parameters
         ]
         x_clone = encoding(x_clone) + circuit
         x_clone.set_parameters(params)
         x_clone = decoding(x_clone)
-        x_clone = torch.as_tensor(backend.to_numpy(x_clone).tolist()).to(x.device)
+        x_clone = torch.as_tensor(
+            backend.to_numpy(x_clone).tolist(), dtype=x.dtype, device=x.device
+        )
         return x_clone
 
     @staticmethod
@@ -166,11 +169,14 @@ def backward(ctx, grad_output: torch.Tensor):
         x_clone = x.clone().detach().cpu().numpy()
         x_clone = ctx.backend.cast(x_clone, dtype=x_clone.dtype)
         params = [
-            ctx.backend.cast(par.clone().detach().cpu().numpy()) for par in parameters
+            ctx.backend.cast(par.clone().detach().cpu().numpy(), dtype=x_clone.dtype)
+            for par in parameters
         ]
         wrt_inputs = not x.is_leaf and ctx.encoding.differentiable
         grad_input, *gradients = (
-            torch.as_tensor(ctx.backend.to_numpy(grad).tolist()).to(x.device)
+            torch.as_tensor(
+                ctx.backend.to_numpy(grad).tolist(), dtype=x.dtype, device=x.device
+            )
             for grad in ctx.differentiation.evaluate(
                 x_clone,
                 ctx.encoding,

src/qiboml/operations/differentiation.py

@@ -131,7 +131,7 @@ def one_parameter_shift(
         generator_eigenval = gate.generator_eigenvalue()
         s = np.pi / (4 * generator_eigenval)
 
-        tmp_params = backend.cast(parameters, copy=True)
+        tmp_params = backend.cast(parameters, copy=True, dtype=parameters[0].dtype)
         tmp_params = self.shift_parameter(tmp_params, parameter_index, s, backend)
 
         circuit.set_parameters(tmp_params)