def _unpack_index(index):
""" Parse index. Always return a tuple of the form (row, col).
Valid type for row/col is integer, slice, or array of integers.
Returns:
resulting row & col indices : single integer, slice, or
array of integers. If row & column indices are supplied
explicitly, they are used as the major/minor indices.
If only one index is supplied, the minor index is
assumed to be all (e.g., [maj, :]).
"""
# First, check if indexing with single boolean matrix.
if ((isinstance(index, (spmatrix, cupy.ndarray, numpy.ndarray))
or _try_is_scipy_spmatrix(index)) and index.ndim == 2
and index.dtype.kind == 'b'):
return index.nonzero()
# Parse any ellipses.
index = _eliminate_ellipsis(index)
# Next, parse the tuple or object
if isinstance(index, tuple):
if len(index) == 2:
row, col = index
elif len(index) == 1:
row, col = index[0], slice(None)
else:
raise IndexError('invalid number of indices')
else:
idx = _compatible_boolean_index(index)
if idx is None:
row, col = index, slice(None)
elif idx.ndim < 2:
return _boolean_index_to_array(idx), slice(None)
elif idx.ndim == 2:
return idx.nonzero()
# Next, check for validity and transform the index as needed.
if isspmatrix(row) or isspmatrix(col):
# Supporting sparse boolean indexing with both row and col does
# not work because spmatrix.ndim is always 2.
raise IndexError('Indexing with sparse matrices is not supported '
'except boolean indexing where matrix and index '
'are equal shapes.')
bool_row = _compatible_boolean_index(row)
bool_col = _compatible_boolean_index(col)
if bool_row is not None:
row = _boolean_index_to_array(bool_row)
if bool_col is not None:
col = _boolean_index_to_array(bool_col)
return row, col
def _add(self, other, lhs_negative, rhs_negative):
if cupy.isscalar(other):
if other == 0:
if lhs_negative:
return -self
else:
return self.copy()
else:
raise NotImplementedError(
'adding a nonzero scalar to a sparse matrix is not '
'supported')
elif base.isspmatrix(other):
alpha = -1 if lhs_negative else 1
beta = -1 if rhs_negative else 1
return self._add_sparse(other, alpha, beta)
elif base.isdense(other):
if lhs_negative:
if rhs_negative:
return -self.todense() - other
else:
return other - self.todense()
else:
if rhs_negative:
return self.todense() - other
else:
return self.todense() + other
else:
return NotImplemented
def __truediv__(self, other):
"""Point-wise division by another matrix, vector or scalar"""
if _util.isscalarlike(other):
dtype = self.dtype
if dtype == numpy.float32:
# Note: This is a work-around to make the output dtype the same
# as SciPy. It might be SciPy version dependent.
dtype = numpy.float64
dtype = cupy.result_type(dtype, other)
d = cupy.reciprocal(other, dtype=dtype)
return multiply_by_scalar(self, d)
elif _util.isdense(other):
other = cupy.atleast_2d(other)
check_shape_for_pointwise_op(self.shape, other.shape)
return self.todense() / other
elif base.isspmatrix(other):
# Note: If broadcasting is needed, an exception is raised here for
# compatibility with SciPy, as SciPy does not support broadcasting
# in the "sparse / sparse" case.
check_shape_for_pointwise_op(self.shape,
other.shape,
allow_broadcasting=False)
dtype = numpy.promote_types(self.dtype, other.dtype)
if dtype.char not in 'FD':
dtype = numpy.promote_types(numpy.float64, dtype)
# Note: The following implementation converts two sparse matrices
# into dense matrices and then performs a point-wise division,
# which can use lots of memory.
self_dense = self.todense().astype(dtype, copy=False)
return self_dense / other.todense()
raise NotImplementedError
def __mul__(self, other):
if cupy.isscalar(other):
self.sum_duplicates()
return self._with_data(self.data * other)
elif isspmatrix_csr(other):
self.sum_duplicates()
other.sum_duplicates()
return cusparse.csrgemm(self, other)
elif csc.isspmatrix_csc(other):
self.sum_duplicates()
other.sum_duplicates()
return cusparse.csrgemm(self, other.T, transb=True)
elif base.isspmatrix(other):
return self * other.tocsr()
elif base.isdense(other):
if other.ndim == 0:
self.sum_duplicates()
return self._with_data(self.data * other)
elif other.ndim == 1:
self.sum_duplicates()
return cusparse.csrmv(self, cupy.asfortranarray(other))
elif other.ndim == 2:
self.sum_duplicates()
return cusparse.csrmm2(self, cupy.asfortranarray(other))
else:
raise ValueError('could not interpret dimensions')
else:
return NotImplemented
def __setitem__(self, key, x):
row, col = self._parse_indices(key)
if isinstance(row, _int_scalar_types) and\
isinstance(col, _int_scalar_types):
x = cupy.asarray(x, dtype=self.dtype)
if x.size != 1:
raise ValueError('Trying to assign a sequence to an item')
self._set_intXint(row, col, x.flat[0])
return
if isinstance(row, slice):
row = cupy.arange(*row.indices(self.shape[0]))[:, None]
else:
row = cupy.atleast_1d(row)
if isinstance(col, slice):
col = cupy.arange(*col.indices(self.shape[1]))[None, :]
if row.ndim == 1:
row = row[:, None]
else:
col = cupy.atleast_1d(col)
i, j = cupy.broadcast_arrays(row, col)
if i.shape != j.shape:
raise IndexError('number of row and column indices differ')
if isspmatrix(x):
if i.ndim == 1:
# Inner indexing, so treat them like row vectors.
i = i[None]
j = j[None]
broadcast_row = x.shape[0] == 1 and i.shape[0] != 1
broadcast_col = x.shape[1] == 1 and i.shape[1] != 1
if not ((broadcast_row or x.shape[0] == i.shape[0]) and
(broadcast_col or x.shape[1] == i.shape[1])):
raise ValueError('shape mismatch in assignment')
if x.size == 0:
return
x = x.tocoo(copy=True)
x.sum_duplicates()
self._set_arrayXarray_sparse(i, j, x)
else:
# Make x and i into the same shape
x = cupy.asarray(x, dtype=self.dtype)
x, _ = cupy.broadcast_arrays(x, i)
if x.size == 0:
return
x = x.reshape(i.shape)
self._set_arrayXarray(i, j, x)
def __mul__(self, other):
if cupy.isscalar(other):
self.sum_duplicates()
return self._with_data(self.data * other)
elif isspmatrix_csr(other):
self.sum_duplicates()
other.sum_duplicates()
return cusparse.csrgemm(self, other)
elif csc.isspmatrix_csc(other):
self.sum_duplicates()
other.sum_duplicates()
return cusparse.csrgemm(self, other.T, transb=True)
elif base.isspmatrix(other):
return self * other.tocsr()
elif base.isdense(other):
if other.ndim == 0:
self.sum_duplicates()
return self._with_data(self.data * other)
elif other.ndim == 1:
self.sum_duplicates()
other = cupy.asfortranarray(other)
# csrmvEx does not work if nnz == 0
if self.nnz > 0 and cusparse.csrmvExIsAligned(self, other):
if cupy.cuda.cub_enabled and other.flags.c_contiguous:
return device_csrmv(self.shape[0], self.shape[1],
self.nnz, self.data, self.indptr,
self.indices, other)
else:
return cusparse.csrmvEx(self, other)
else:
return cusparse.csrmv(self, other)
elif other.ndim == 2:
self.sum_duplicates()
return cusparse.csrmm2(self, cupy.asfortranarray(other))
else:
raise ValueError('could not interpret dimensions')
else:
return NotImplemented
def __mul__(self, other):
if cupy.isscalar(other):
self.sum_duplicates()
return self._with_data(self.data * other)
elif isspmatrix_csr(other):
self.sum_duplicates()
other.sum_duplicates()
if cusparse.check_availability('csrgemm2'):
return cusparse.csrgemm2(self, other)
elif cusparse.check_availability('csrgemm'):
return cusparse.csrgemm(self, other)
else:
raise NotImplementedError
elif csc.isspmatrix_csc(other):
self.sum_duplicates()
other.sum_duplicates()
if cusparse.check_availability('csrgemm') and not runtime.is_hip:
# trans=True is still buggy as of ROCm 4.2.0
return cusparse.csrgemm(self, other.T, transb=True)
elif cusparse.check_availability('csrgemm2'):
b = other.tocsr()
b.sum_duplicates()
return cusparse.csrgemm2(self, b)
else:
raise NotImplementedError
elif base.isspmatrix(other):
return self * other.tocsr()
elif base.isdense(other):
if other.ndim == 0:
self.sum_duplicates()
return self._with_data(self.data * other)
elif other.ndim == 1:
self.sum_duplicates()
other = cupy.asfortranarray(other)
# need extra padding to ensure not stepping on the CUB bug,
# see cupy/cupy#3679 for discussion
is_cub_safe = (self.indptr.data.mem.size >
self.indptr.size * self.indptr.dtype.itemsize)
# CUB spmv is buggy since CUDA 11.0, see
# https://github.com/cupy/cupy/issues/3822#issuecomment-782607637
is_cub_safe &= (cub._get_cuda_build_version() < 11000)
for accelerator in _accelerator.get_routine_accelerators():
if (accelerator == _accelerator.ACCELERATOR_CUB
and not runtime.is_hip and is_cub_safe
and other.flags.c_contiguous):
return cub.device_csrmv(self.shape[0], self.shape[1],
self.nnz, self.data,
self.indptr, self.indices,
other)
if (cusparse.check_availability('csrmvEx') and self.nnz > 0
and cusparse.csrmvExIsAligned(self, other)):
# csrmvEx does not work if nnz == 0
csrmv = cusparse.csrmvEx
elif cusparse.check_availability('csrmv'):
csrmv = cusparse.csrmv
elif cusparse.check_availability('spmv'):
csrmv = cusparse.spmv
else:
raise NotImplementedError
return csrmv(self, other)
elif other.ndim == 2:
self.sum_duplicates()
if cusparse.check_availability('csrmm2'):
csrmm = cusparse.csrmm2
elif cusparse.check_availability('spmm'):
csrmm = cusparse.spmm
else:
raise NotImplementedError
return csrmm(self, cupy.asfortranarray(other))
else:
raise ValueError('could not interpret dimensions')
else:
return NotImplemented
def __mul__(self, other):
if cupy.isscalar(other):
self.sum_duplicates()
return self._with_data(self.data * other)
elif isspmatrix_csr(other):
self.sum_duplicates()
other.sum_duplicates()
if cusparse.check_availability('csrgemm2'):
return cusparse.csrgemm2(self, other)
elif cusparse.check_availability('csrgemm'):
return cusparse.csrgemm(self, other)
else:
raise NotImplementedError
elif csc.isspmatrix_csc(other):
self.sum_duplicates()
other.sum_duplicates()
if cusparse.check_availability('csrgemm'):
return cusparse.csrgemm(self, other.T, transb=True)
elif cusparse.check_availability('csrgemm2'):
b = other.tocsr()
b.sum_duplicates()
return cusparse.csrgemm2(self, b)
else:
raise NotImplementedError
elif base.isspmatrix(other):
return self * other.tocsr()
elif base.isdense(other):
if other.ndim == 0:
self.sum_duplicates()
return self._with_data(self.data * other)
elif other.ndim == 1:
self.sum_duplicates()
other = cupy.asfortranarray(other)
# csrmvEx does not work if nnz == 0
if self.nnz > 0 and cusparse.csrmvExIsAligned(self, other):
for accelerator in _accelerator.get_routine_accelerators():
if (accelerator == _accelerator.ACCELERATOR_CUB
and other.flags.c_contiguous):
return cub.device_csrmv(self.shape[0],
self.shape[1], self.nnz,
self.data, self.indptr,
self.indices, other)
return cusparse.csrmvEx(self, other)
else:
if cusparse.check_availability('csrmv'):
csrmv = cusparse.csrmv
elif cusparse.check_availability('spmv'):
csrmv = cusparse.spmv
else:
raise NotImplementedError
return csrmv(self, other)
elif other.ndim == 2:
self.sum_duplicates()
if cusparse.check_availability('csrmm2'):
csrmm = cusparse.csrmm2
elif cusparse.check_availability('spmm'):
csrmm = cusparse.spmm
else:
raise NotImplementedError
return csrmm(self, cupy.asfortranarray(other))
else:
raise ValueError('could not interpret dimensions')
else:
return NotImplemented
