jax.experimental.sparse.coo_fromdense

jax.experimental.sparse.coo_fromdense#

jax.experimental.sparse.coo_fromdense(mat, *, nse=None, index_dtype=<class 'jax.numpy.int32'>)[source]#

Create a COO-format sparse matrix from a dense matrix.

Parameters:

mat (Array) – array to be converted to COO.
nse (int | None | None) – number of specified entries in mat. If not specified, it will be computed from the input matrix.
index_dtype (DTypeLike) – dtype of sparse indices

Returns:

COO representation of the matrix.

Return type:

mat_coo