jax.experimental.sparse.JAXSparse

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

Base class for high-level JAX sparse objects.

Parameters:

• args (tuple[Array, ...])

• shape (tuple[int, ...])

__init__(args, *, shape)

Parameters:

• args (tuple[Array, ...])

• shape (Sequence[int])

Methods

__init__(args, *, shape)
block_until_ready()
sum(*args, **kwargs)
transpose([axes])
tree_flatten()
tree_unflatten(aux_data, children)

Attributes

T
ndim
size
data
shape
nse
dtype