keras/src/backend/torch/math.py/_segment_reduction_fn
def _segment_reduction_fn(data, segment_ids, reduction_method, num_segments):
    num_repeats = torch.prod(
        torch.tensor(data.shape[1:], device=get_device())
    ).long()
    # To use `scatter_add` in torch, we need to replicate `segment_ids` into the
    # shape of `data`.
    segment_ids = (
        segment_ids.repeat_interleave(num_repeats)
        .view(*data.shape)
        .type(torch.int64)
    )
    num_segments = num_segments or len(torch.unique(segment_ids))

    # .scatter_add does not support -1 in the indices.
    # Add all out-of-bound indices value to an extra dimension after
    # num_segments, which is removed before returning the result.

    # Replacing the out-of-bound indices.
    segment_ids = torch.where(segment_ids >= 0, segment_ids, num_segments)
    segment_ids = torch.where(
        segment_ids < num_segments, segment_ids, num_segments
    )

    # Add one more dimension to the result shape with the "+1".
    shape = (num_segments + 1,) + tuple(data.shape[1:])

    if reduction_method == "amax":
        result = torch.ones(*shape, device=get_device()) * -float("Inf")
    else:
        result = torch.zeros(*shape, device=get_device())

    result = result.scatter_reduce(
        0, segment_ids, data.float(), reduction_method
    )

    # Removing the extra dimension.
    result = result[:-1, ...]

    return result.type(data.dtype)