keras/src/ops/math.py/SegmentReduction/__init__
class SegmentReduction:    def __init__(self, num_segments=None, sorted=False):
        super().__init__()
        self.num_segments = num_segments
        self.sorted = sorted