keras/src/backend/jax/numpy.py/sum
def sum(x, axis=None, keepdims=False):
    x = convert_to_tensor(x)
    if isinstance(x, jax_sparse.BCOO):
        if axis is None:
            axis = tuple(range(len(x.shape)))
        (
            canonical_axis,
            keep_dims_shape,
            broadcast_dimensions,
        ) = sparse.axis_shape_dims_for_broadcast_in_dim(
            axis, x.shape, insert_dims=False
        )
        output = jax_sparse.bcoo_reduce_sum(x, axes=canonical_axis)
        if keepdims:
            # `bcoo_reduce_sum` does not support keepdims, neither does
            # sparsify(jnp.sum), so we recreate the empty dimensions.
            output = jax_sparse.bcoo_broadcast_in_dim(
                output,
                shape=keep_dims_shape,
                broadcast_dimensions=broadcast_dimensions,
            )
        return output
    return jnp.sum(x, axis=axis, keepdims=keepdims)