fix: adjustments for cupy

src/qibojit/backends/clifford_operations_cpu.py

@@ -260,9 +260,9 @@ def CY(symplectic_matrix, control_q, target_q, nqubits):
     cache=True,
     fastmath=True,
 )
-def _rowsum(symplectic_matrix, h, i, dim, determined=False):
-    xi, zi = symplectic_matrix[i, :dim], symplectic_matrix[i, dim:-1]
-    xh, zh = symplectic_matrix[h, :dim], symplectic_matrix[h, dim:-1]
+def _rowsum(symplectic_matrix, h, i, nqubits, determined=False):
+    xi, zi = symplectic_matrix[i, :nqubits], symplectic_matrix[i, nqubits:-1]
+    xh, zh = symplectic_matrix[h, :nqubits], symplectic_matrix[h, nqubits:-1]
     if determined:
         g_r = np.zeros(h.shape[0], dtype=np.uint8)
         g_xi_xh = xi.copy()
@@ -286,11 +286,11 @@ def _rowsum(symplectic_matrix, h, i, dim, determined=False):
             g_zi_zh[j] = zi_zh
         else:
             symplectic_matrix[h[j], -1] = r
-            symplectic_matrix[h[j], :dim] = xi_xh
-            symplectic_matrix[h[j], dim:-1] = zi_zh
+            symplectic_matrix[h[j], :nqubits] = xi_xh
+            symplectic_matrix[h[j], nqubits:-1] = zi_zh
     if determined:
         for j in prange(len(g_r)):  # pylint: disable=not-an-iterable
             symplectic_matrix[h[0], -1] ^= g_r[j]
-            symplectic_matrix[h[0], :dim] ^= g_xi_xh[j]
-            symplectic_matrix[h[0], dim:-1] ^= g_zi_zh[j]
+            symplectic_matrix[h[0], :nqubits] ^= g_xi_xh[j]
+            symplectic_matrix[h[0], nqubits:-1] ^= g_zi_zh[j]
     return symplectic_matrix

src/qibojit/backends/clifford_operations_gpu.py

@@ -657,66 +657,57 @@ def CY(symplectic_matrix, control_q, target_q, nqubits):
 
 
 def _rowsum(symplectic_matrix, h, i, nqubits, determined=False):
-    dim = _get_dim(nqubits)
     nrows = len(h)
     exp = cp.zeros(len(h), dtype=int)
+    packed_nqubits = _get_packed_size(nqubits)
+    row_dim = _get_dim(packed_nqubits)
     apply_rowsum(
         GRIDDIM_2D,
         (BLOCKDIM,),
-        (symplectic_matrix, h, i, nqubits, determined, nrows, exp, dim),
+        (symplectic_matrix, h, i, packed_nqubits, determined, nrows, exp, row_dim),
     )
     return symplectic_matrix
 
 
-def _get_p(state, q, nqubits):
-    dim = _get_dim(nqubits)
-    return state.reshape(dim, dim)[nqubits:-1, q].nonzero()[0]
-
-
 def _random_outcome(state, p, q, nqubits):
     dim = _get_dim(nqubits)
     p = p[0] + nqubits
-    idx_pq = p * dim + q
-    tmp = state[idx_pq].copy()
-    state[idx_pq] = False
-    h = state.reshape(dim, dim)[:-1, q].nonzero()[0]
-    state[idx_pq] = tmp
-    packed_dim = _get_packed_dim(state.shape[1])
+    tmp = state[p, q].copy()
+    state[p, q] = 0
+    h = state[:-1, q].nonzero()[0]
+    state[p, q] = tmp
     if h.shape[0] > 0:
         state = _pack_for_measurements(state, nqubits)
         state = _rowsum(
-            state,
+            state.ravel(),
             h,
             p.astype(cp.uint) * cp.ones(h.shape[0], dtype=np.uint),
-            packed_dim,
+            _get_packed_size(nqubits),
             False,
         )
-        state = _unpack_for_measurements(state, nqubits)
-    state = state.reshape(dim, dim)
+        state = _unpack_for_measurements(state.reshape(dim, dim), nqubits)
     state[p - nqubits, :] = state[p, :]
     outcome = cp.random.randint(2, size=None, dtype=cp.uint)
-    state[p, :] = False
+    state[p, :] = 0
     state[p, -1] = outcome.astype(cp.uint8)
-    state[p, nqubits + q] = True
-    return state.ravel(), outcome
+    state[p, nqubits + q] = 1
+    return state, outcome
 
 
 def _determined_outcome(state, q, nqubits):
+    state[-1, :] = 0
+    idx = (state[:nqubits, q].nonzero()[0] + nqubits).astype(np.uint)
     dim = _get_dim(nqubits)
-    state = state.reshape(dim, dim)
-    state[-1, :] = False
-    idx = state[:nqubits, q].nonzero()[0] + nqubits
-    packed_dim = _get_packed_dim(state.shape[1])
     state = _pack_for_measurements(state, nqubits)
     state = _rowsum(
         state.ravel(),
         (2 * nqubits * cp.ones(idx.shape, dtype=np.uint)).astype(np.uint),
         idx.astype(np.uint),
-        packed_dim,
+        _get_packed_size(nqubits),
         True,
     )
-    state = _unpack_for_measurements(state, nqubits)
-    return state, state[dim * dim - 1].astype(cp.uint)
+    state = _unpack_for_measurements(state.reshape(dim, dim), nqubits)
+    return state, state[-1, -1]
 
 
 def _packbits(array, axis):
@@ -728,6 +719,14 @@ def _unpackbits(array, axis):
     return cp.array(numpy.unpackbits(array.get(), axis=axis), dtype=cp.uint8)
 
 
+def _init_state_for_measurements(state, nqubits, collapse):
+    dim = _get_dim(nqubits)
+    if collapse:
+        return _unpackbits(state.reshape(dim, dim), axis=0)[:dim]
+    else:
+        return state.copy()
+
+
 def cast(x, dtype=None, copy=False):
     if dtype is None:
         dtype = "complex128"