jax.ops.segment_sum(data, segment_ids, num_segments=None, indices_are_sorted=False, unique_indices=False, bucket_size=None)[source]

Computes the sum within segments of an array.

Similar to TensorFlow’s segment_sum: https://www.tensorflow.org/api_docs/python/tf/math/segment_sum

  • data – an array with the values to be summed.

  • segment_ids – an array with integer dtype that indicates the segments of data (along its leading axis) to be summed. Values can be repeated and need not be sorted. Values outside of the range [0, num_segments) are wrapped into that range by applying jnp.mod.

  • num_segments – optional, an int with positive value indicating the number of segments. The default is set to be the minimum number of segments that would support all positive and negative indices in segment_ids calculated as max(max(segment_ids) + 1, max(-segment_ids)). Since num_segments determines the size of the output, a static value must be provided to use segment_sum in a jit-compiled function.

  • indices_are_sorted – whether segment_ids is known to be sorted.

  • unique_indices – whether segment_ids is known to be free of duplicates.

  • bucket_size – size of bucket to group indices into. segment_sum is performed on each bucket separately to improve numerical stability of addition. Default None means no bucketing.


An array with shape (num_segments,) + data.shape[1:] representing the segment sums.