def segment_sum(data, segment_ids, num_segments=None, sorted=False):
    """Computes the sum of segments in a tensor.

    Args:
        data: Input tensor.
        segment_ids: A N-D tensor containing segment indices for each
            element in `data`. Num dims for segment ids should be strictly
            smaller or equal to number of dims in data.
        num_segments: An integer representing the total number of
            segments. If not specified, it is inferred from the maximum
            value in `segment_ids`.
        sorted: A boolean indicating whether `segment_ids` is sorted.
            Defaults to `False`.

    Returns:
        A tensor containing the sum of segments, where each element
        represents the sum of the corresponding segment in `data`.

    Example:

    >>> data = keras.ops.convert_to_tensor([1, 2, 10, 20, 100, 200])
    >>> segment_ids = keras.ops.convert_to_tensor([0, 0, 1, 1, 2, 2])
    >>> num_segments = 3
    >>> keras.ops.segment_sum(data, segment_ids,num_segments)
    array([3, 30, 300], dtype=int32)
    """
    _segment_reduce_validation(data, segment_ids)
    if any_symbolic_tensors((data,)):
        return SegmentSum(num_segments, sorted).symbolic_call(data, segment_ids)
    return backend.math.segment_sum(
        data, segment_ids, num_segments=num_segments, sorted=sorted
    )
Computes the sum of segments in a tensor.

    Args:
        data: Input tensor.
        segment_ids: A N-D tensor containing segment indices for each
            element in `data`. Num dims for segment ids should be strictly
            smaller or equal to number of dims in data.
        num_segments: An integer representing the total number of
            segments. If not specified, it is inferred from the maximum
            value in `segment_ids`.
        sorted: A boolean indicating whether `segment_ids` is sorted.
            Defaults to `False`.

    Returns:
        A tensor containing the sum of segments, where each element
        represents the sum of the corresponding segment in `data`.

    Example:

    >>> data = keras.ops.convert_to_tensor([1, 2, 10, 20, 100, 200])
    >>> segment_ids = keras.ops.convert_to_tensor([0, 0, 1, 1, 2, 2])
    >>> num_segments = 3
    >>> keras.ops.segment_sum(data, segment_ids,num_segments)
    array([3, 30, 300], dtype=int32)
    
keras/src/ops/math.py