jax.experimental.sparse.CSR

class jax.experimental.sparse.CSR(args, *, shape)

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.

Parameters:

shape (tuple[int, int])

__init__(args, *, shape)

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