keras/src/backend/numpy/math.py/_segment_reduction_fn
def _segment_reduction_fn(
    data, segment_ids, reduction_method, num_segments, sorted
):
    if num_segments is None:
        num_segments = np.amax(segment_ids) + 1

    valid_indices = segment_ids >= 0  # Ignore segment_ids that are -1
    valid_data = data[valid_indices]
    valid_segment_ids = segment_ids[valid_indices]

    data_shape = list(valid_data.shape)
    data_shape[0] = (
        num_segments  # Replace first dimension (which corresponds to segments)
    )

    if reduction_method == np.maximum:
        result = np.ones(data_shape, dtype=valid_data.dtype) * -np.inf
    else:
        result = np.zeros(data_shape, dtype=valid_data.dtype)

    if sorted:
        reduction_method.at(result, valid_segment_ids, valid_data)
    else:
        sort_indices = np.argsort(valid_segment_ids)
        sorted_segment_ids = valid_segment_ids[sort_indices]
        sorted_data = valid_data[sort_indices]

        reduction_method.at(result, sorted_segment_ids, sorted_data)

    return result