jax.ops.segment\_sum ==================== .. currentmodule:: jax.ops .. autofunction:: segment_sum