jax.experimental.sparse.COO

class jax.experimental.sparse.COO(args, *, shape, rows_sorted=False, cols_sorted=False)

Experimental COO 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:

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

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

• rows_sorted (bool)

• cols_sorted (bool)

__init__(args, *, shape, rows_sorted=False, cols_sorted=False)

Parameters:

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

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

• rows_sorted (bool)

• cols_sorted (bool)

Methods

__init__(args, *, shape[, rows_sorted, ...])	param args:
block_until_ready()
fromdense(mat, *[, nse, index_dtype])	param mat:
sum(*args, **kwargs)
todense()	rtype: Array
transpose([axes])	param axes:
tree_flatten()	rtype: tuple[tuple[Array, Array, Array], dict[str, Any]]
tree_unflatten(aux_data, children)

Attributes

T
dtype
ndim
nse
size
data
row
col
shape