jax.experimental.sparse.bcoo_update_layout#

jax.experimental.sparse.bcoo_update_layout(mat, *, n_batch=None, n_dense=None, on_inefficient='error')[source]#

Update the storage layout (i.e. n_batch & n_dense) of a BCOO matrix.

In many cases this can be done without introducing undue storage overhead. However, increasing mat.n_batch or mat.n_dense will lead to very inefficient storage, with many explicitly-stored zeros, unless the new batch or dense dimensions have size 0 or 1. In such cases, bcoo_update_layout will raise a SparseEfficiencyError. This can be silenced by specifying the on_inefficient argument.

Parameters:

  • mat (BCOO) – BCOO array

  • n_batch (int | None) – optional(int) the number of batch dimensions in the output matrix. If None, then n_batch = mat.n_batch.

  • n_dense (int | None) – optional(int) the number of dense dimensions in the output matrix. If None, then n_dense = mat.n_dense.

  • on_inefficient (str | None) – optional(string), one of ['error', 'warn', None]. Specify the behavior in case of an inefficient reconfiguration. This is defined as a reconfiguration where the size of the resulting representation is much larger than the size of the input representation.

Returns:

BCOO array

A BCOO array representing the same sparse array as the input, with the specified layout. mat_out.todense() will match mat.todense() up to appropriate precision.

Return type:

mat_out