keras/src/ops/math.py/SegmentMax/call
class SegmentMax:    def call(self, data, segment_ids):
        _segment_reduce_validation(data, segment_ids)
        return backend.math.segment_max(
            data,
            segment_ids,
            num_segments=self.num_segments,
            sorted=self.sorted,
        )