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.

Return type

BCOO

Returns

A BCOO-format array containing the result.