keras/src/ops/math.py/SegmentReduction/compute_output_spec
class SegmentReduction:    def compute_output_spec(self, data, _):
        output_shape = (self.num_segments,) + tuple(data.shape[1:])
        return KerasTensor(shape=output_shape, dtype=data.dtype)