jax.experimental.sparse.bcoo_extract

jax.experimental.sparse.bcoo_extract(sparr, arr, *, assume_unique=None)

Extract values from a dense array according to the sparse array’s indices.

Parameters:

• sparr (BCOO) – BCOO array whose indices will be used for the output.

• arr (ArrayLike) – ArrayLike with shape equal to self.shape

• assume_unique (bool | None) – bool, defaults to sparr.unique_indices If True, extract values for every index, even if index contains duplicates. If False, duplicate indices will have their values summed and returned in the position of the first index.

Returns:

a BCOO array with the same sparsity pattern as self.

Return type:

extracted