jax.experimental.sparse.bcoo_reduce_sum

jax.experimental.sparse.bcoo_reduce_sum#

jax.experimental.sparse.bcoo_reduce_sum(mat, *, axes)[source]#

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.

Returns:

A BCOO-format array containing the result.

Return type:

BCOO

Contents