jax.experimental.sparse.bcoo_sum_duplicates

jax.experimental.sparse.bcoo_sum_duplicates(mat, nse=None)

Sums duplicate indices within a BCOO array, returning an array with sorted indices.

Parameters

• mat (BCOO) – BCOO array

• nse (Optional[int]) – integer (optional). The number of specified elements in the output matrix. This must be specified for bcoo_sum_duplicates to be compatible with JIT and other JAX transformations. If not specified, the optimal nse will be computed based on the contents of the data and index arrays. If specified nse is larger than necessary, data and index arrays will be padded with standard fill values. If smaller than necessary, data elements will be dropped from the output matrix.

Returns

BCOO array with sorted indices and no duplicate indices.

Return type

mat_out