brainpy.math.csr_matvec#

brainpy.math.csr_matvec(data, indices, indptr, vector, *, shape, method='vector')[source]#

CSR sparse matrix product with a dense vector, which outperforms the cuSPARSE algorithm.

This function supports JAX transformations, including jit(), 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 dtype indices.dtype.

  • vector (ndarray) – An array of shape (shape[0] if transpose else shape[1],) and dtype data.dtype.

  • shape (tuple of int) – A length-2 tuple representing the matrix shape.

  • method (str) – The computing method used in GPU backend. Currently, we support scalar, vector and adaptive.

Returns:

y – The array of shape (shape[1] if transpose else shape[0],) representing the matrix vector product.

Return type:

ndarry