keras/src/backend/tensorflow/math.py/segment_sum
def segment_sum(data, segment_ids, num_segments=None, sorted=False):
    if sorted:
        if num_segments is not None:
            raise ValueError(
                "Argument `num_segments` cannot be set when sorted is True "
                "when using the tensorflow backend."
                f"Received: num_segments={num_segments}, sorted={sorted}."
            )
        return tf.math.segment_sum(data, segment_ids)
    else:
        if num_segments is None:
            unique_segment_ids, _ = tf.unique(segment_ids)
            num_segments = tf.shape(unique_segment_ids)[0]
        return tf.math.unsorted_segment_sum(data, segment_ids, num_segments)