jax.experimental.sparse.bcoo_reduce_sum

jax.experimental.sparse.bcoo_reduce_sum(mat, *, axes)

Sum array element over given axes.

Parameters

• mat (BCOO) – A BCOO-format array.

• shape – The shape of the target array.

• axes (Sequence[int]) – A tuple or list or ndarray which contains axes of mat over which sum is performed.

Return type

BCOO

Returns

A BCOO-format array containing the result.