jax.experimental.sparse.CSR

Contents

jax.experimental.sparse.CSR#

class jax.experimental.sparse.CSR(args, *, shape)[source]#

Experimental CSR matrix implemented in JAX.

Note: this class has minimal compatibility with JAX transforms such as grad and autodiff, and offers very little functionality. In general you should prefer jax.experimental.sparse.BCOO.

Additionally, there are known failures in the case that nse is larger than the true number of nonzeros in the represented matrix. This situation is better handled in BCOO.

__init__(args, *, shape)[source]#

Methods

__init__(args, *, shape)

block_until_ready()

fromdense(mat, *[, nse, index_dtype])

sum(*args, **kwargs)

todense()

transpose([axes])

tree_flatten()

tree_unflatten(aux_data, children)

Attributes

T

dtype

ndim

nse

size

data

indices

indptr

shape