csrmv

Contents

csrmv#

class brainpy.math.sparse.csrmv(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 (Union[float, Array, Array]) – An array of shape (nse,).

  • indices (Union[Array, Array]) – An array of shape (nse,).

  • indptr (Union[Array, Array]) – An array of shape (shape[0] + 1,) and dtype indices.dtype.

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

  • shape (Tuple[int, 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

Contents