brainpy.math.cusparse_coo_matvec#

brainpy.math.cusparse_coo_matvec(data, row, col, vector, *, shape, rows_sorted=False, cols_sorted=False, transpose=False)[source]#

Product of COO 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,).

  • row (ndarray) – An array of shape (nse,).

  • col (ndarray) – An array of shape (nse,) and dtype row.dtype.

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

  • shape (tuple of int) – The shape of the sparse matrix.

  • rows_sorted (bool) – Row index are sorted.

  • cols_sorted (bool) – Column index are sorted.

  • transpose (bool) – A boolean specifying whether to transpose the sparse matrix before computing.

Returns:

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

Return type:

ndarray