def csrgeam(a, b, alpha=1, beta=1): """Matrix-matrix addition. .. math:: C = \\alpha A + \\beta B Args: a (cupyx.scipy.sparse.csr_matrix): Sparse matrix A. b (cupyx.scipy.sparse.csr_matrix): Sparse matrix B. alpha (float): Coefficient for A. beta (float): Coefficient for B. Returns: cupyx.scipy.sparse.csr_matrix: Result matrix. """ if not check_availability('csrgeam'): raise RuntimeError('csrgeam is not available.') if not isinstance(a, cupyx.scipy.sparse.csr_matrix): raise TypeError('unsupported type (actual: {})'.format(type(a))) if not isinstance(b, cupyx.scipy.sparse.csr_matrix): raise TypeError('unsupported type (actual: {})'.format(type(b))) assert a.has_canonical_format assert b.has_canonical_format if a.shape != b.shape: raise ValueError('inconsistent shapes') handle = device.get_cusparse_handle() m, n = a.shape a, b = _cast_common_type(a, b) nnz = numpy.empty((), 'i') cusparse.setPointerMode(handle, cusparse.CUSPARSE_POINTER_MODE_HOST) c_descr = MatDescriptor.create() c_indptr = cupy.empty(m + 1, 'i') cusparse.xcsrgeamNnz(handle, m, n, a._descr.descriptor, a.nnz, a.indptr.data.ptr, a.indices.data.ptr, b._descr.descriptor, b.nnz, b.indptr.data.ptr, b.indices.data.ptr, c_descr.descriptor, c_indptr.data.ptr, nnz.ctypes.data) c_indices = cupy.empty(int(nnz), 'i') c_data = cupy.empty(int(nnz), a.dtype) alpha = numpy.array(alpha, a.dtype).ctypes beta = numpy.array(beta, a.dtype).ctypes _call_cusparse('csrgeam', a.dtype, handle, m, n, alpha.data, a._descr.descriptor, a.nnz, a.data.data.ptr, a.indptr.data.ptr, a.indices.data.ptr, beta.data, b._descr.descriptor, b.nnz, b.data.data.ptr, b.indptr.data.ptr, b.indices.data.ptr, c_descr.descriptor, c_data.data.ptr, c_indptr.data.ptr, c_indices.data.ptr) c = cupyx.scipy.sparse.csr_matrix((c_data, c_indices, c_indptr), shape=a.shape) c._has_canonical_format = True return c
def csrgeam(a, b, alpha=1, beta=1): """Matrix-matrix addition. .. math:: C = \\alpha A + \\beta B Args: a (cupy.sparse.csr_matrix): Sparse matrix A. b (cupy.sparse.csr_matrix): Sparse matrix B. alpha (float): Coefficient for A. beta (float): Coefficient for B. Returns: cupy.sparse.csr_matrix: Result matrix. """ if a.shape != b.shape: raise ValueError('inconsistent shapes') handle = device.get_cusparse_handle() m, n = a.shape a, b = _cast_common_type(a, b) nnz = numpy.empty((), 'i') cusparse.setPointerMode( handle, cusparse.CUSPARSE_POINTER_MODE_HOST) c_descr = MatDescriptor.create() c_indptr = cupy.empty(m + 1, 'i') cusparse.xcsrgeamNnz( handle, m, n, a._descr.descriptor, a.nnz, a.indptr.data.ptr, a.indices.data.ptr, b._descr.descriptor, b.nnz, b.indptr.data.ptr, b.indices.data.ptr, c_descr.descriptor, c_indptr.data.ptr, nnz.ctypes.data) c_indices = cupy.empty(int(nnz), 'i') c_data = cupy.empty(int(nnz), a.dtype) alpha = numpy.array(alpha, a.dtype).ctypes beta = numpy.array(beta, a.dtype).ctypes _call_cusparse( 'csrgeam', a.dtype, handle, m, n, alpha.data, a._descr.descriptor, a.nnz, a.data.data.ptr, a.indptr.data.ptr, a.indices.data.ptr, beta.data, b._descr.descriptor, b.nnz, b.data.data.ptr, b.indptr.data.ptr, b.indices.data.ptr, c_descr.descriptor, c_data.data.ptr, c_indptr.data.ptr, c_indices.data.ptr) return cupy.sparse.csr_matrix((c_data, c_indices, c_indptr), shape=a.shape)