brainpy.math.cusparse_csr_matvec#
- brainpy.math.cusparse_csr_matvec(data, indices, indptr, vector, *, shape, transpose=False)[source]#
Product of CSR sparse matrix and a dense vector using cuSPARSE algorithm.
This function supports JAX transformations, including jit(), grad(), vmap() and pmap().
- Parameters:
data (ndarray, float) – An array of shape
(nse,)
.indices (ndarray) – An array of shape
(nse,)
.indptr (ndarray) – An array of shape
(shape[0] + 1,)
and dtypeindices.dtype
.vector (ndarray) – An array of shape
(shape[0] if transpose else shape[1],)
and dtypedata.dtype
.shape (tuple of int) – A length-2 tuple representing the matrix shape.
transpose (bool) – A boolean specifying whether to transpose the sparse matrix before computing.
- Returns:
y – The array of shape
(shape[1] if transpose else shape[0],)
representing the matrix vector product.- Return type:
ndarry