def _major_slice(self, idx, copy=False): """Index along the major axis where idx is a slice object. """ if idx == slice(None): return self.copy() if copy else self M, N = self._swap(*self.shape) start, stop, step = idx.indices(M) M = len(range(start, stop, step)) new_shape = self._swap(M, N) if M == 0 or self.nnz == 0: return self.__class__(new_shape) if step == 1: indptr, indices, data = _index._get_csr_submatrix_major_axis( self.indptr, self.indices, self.data, start, stop) else: rows = cupy.arange(start, start + M * step, step, dtype=self.indptr.dtype) indptr, indices, data = _index._csr_row_index( rows, self.indptr, self.indices, self.data) return self.__class__((data, indices, indptr), shape=new_shape, copy=False)
def _major_index_fancy(self, idx): """Index along the major axis where idx is an array of ints. """ _, N = self._swap(*self.shape) M = idx.size new_shape = self._swap(M, N) if self.nnz == 0 or M == 0: return self.__class__(new_shape) return self.__class__( _index._csr_row_index(self.data, self.indices, self.indptr, idx), shape=new_shape, copy=False)
def _major_index_fancy(self, idx): """Index along the major axis where idx is an array of ints. """ _, N = self._swap(*self.shape) M = len(idx) new_shape = self._swap(M, N) if M == 0: return self.__class__(new_shape) row_nnz = cupy.diff(self.indptr) idx_dtype = self.indices.dtype res_indptr = cupy.zeros(M + 1, dtype=idx_dtype) cupy.cumsum(row_nnz[idx], out=res_indptr[1:]) res_indices, res_data = _index._csr_row_index(idx, self.indptr, self.indices, self.data, res_indptr) return self.__class__((res_data, res_indices, res_indptr), shape=new_shape, copy=False)
def _major_slice(self, idx, copy=False): """Index along the major axis where idx is a slice object. """ if idx == slice(None): return self.copy() if copy else self M, N = self._swap(*self.shape) start, stop, step = idx.indices(M) M = len(range(start, stop, step)) new_shape = self._swap(M, N) if M == 0: return self.__class__(new_shape) row_nnz = cupy.diff(self.indptr) idx_dtype = self.indices.dtype res_indptr = cupy.zeros(M + 1, dtype=idx_dtype) cupy.cumsum(row_nnz[idx], out=res_indptr[1:]) if step == 1: idx_start = self.indptr[start] idx_stop = self.indptr[stop] res_indices = cupy.array(self.indices[idx_start:idx_stop], copy=copy) res_data = cupy.array(self.data[idx_start:idx_stop], copy=copy) else: rows = cupy.arange(start, start + (res_indptr.size - 1) * step, step, dtype=res_indptr.dtype) res_indices, res_data = _index._csr_row_index( rows, self.indptr, self.indices, self.data, res_indptr) return self.__class__((res_data, res_indices, res_indptr), shape=new_shape, copy=False)