keras/src/backend/numpy/math.py/segment_sum
def segment_sum(data, segment_ids, num_segments=None, sorted=False):
    return _segment_reduction_fn(
        data, segment_ids, np.add, num_segments, sorted
    )