Repository: pytorch
Issue: Scatter min/max reduce operation that returns the corresponding indices
Labels: triaged,enhancement,module: scatter & gather ops

### 🚀 The feature, motivation and pitch

There are cases where you may be interested in the maximum index associated with a scatter reduce operation. For example, if you have a tensor of vectors and a corresponding vector of group IDs and want to find the vector with the maximum norm from each group.

For example say a tensor of vectors `x` of shape `[100, 512]` and a corresponding list of group ids in the range `0` to `10`.

Example data

```python
x = torch.randn(100, 512)
group_ids = torch.randint(10, (100,))
```

and then say you want to find the vector from each group with maximum norm.

Using `scatter_reduce` with the `reduce='amax'` argument you can find the maximum norm, but there's currently no built-in way to get the index of the maximum element so that you can find the corresponding vector with that norm.

I'm proposing a `scatter_argreduce` function that would provide this functionality.

```python
x_norm = torch.norm(x, dim=1)
max_idx = torch.scatter_argreduce(x_norm, 0, group_ids, reduce='max')

# y is a [10, 512] tensor containing rows from x with maximum norm for each of the 10 groups
y = x[max_idx]
```

The following diagram gives a graphical representation of what I'm proposing

![Untitled Diagram](https://user-images.githubusercontent.com/194276/186423685-b9f8e607-b0f4-4e4b-87fa-e109f4a7d55f.png)

### Proposed function:

`torch.scatter_argreduce(input, dim, index, reduce, *, output_size=None, default_index=0, return_values=False) -> Tensor`

Parameters

 - **input** (Tensor) - The input tensor
 - **dim** (int) - the axis along which to index
 - **index** (LongTensor) - the indices of elements to scatter and reduce
 - **reduce** (str) - "max" or "min"
 - **output_size** (int, optional) - the size of the output at dimension `dim`.
 - **default_index** (int, optional) - the default index to place at intermediate indices not present in `index`.
 - **return_values** (bool, optional) - additionally return the min/max value (equivalent to gather_reduce output)

A naive python implementation may look something like this (borrowed heavily from [scatter_reduce_two_cpu](https://github.com/pytorch/pytorch/blob/bc2c6edaf163b1a1330e37a6e34caf8c553e4755/aten/src/ATen/native/TensorAdvancedIndexing.cpp#L1274))

```python
def scatter_argreduce(input, dim, index, reduce, *, output_size=None, default_index=0, return_values=False):
    assert -input.dim() <= dim < input.dim(), f"Expected `dim` to be in range {-input.dim()}, to {input.dim()}" \
            f" (got {dim})"
    assert reduce == 'max' or reduce == 'min', f"`reduce` argument must be one of ('min`, 'max'), (got '{reduce}')"
    assert input.size() == index.size(), f"Shape mismatched between `input` (got {input.size()} and `index`" \
            f" (got {index.size()})"
    assert output_size is None or (isinstance(output_size, int) and output_size >= 0), "Invalid output size," \
            f"should be non-negative int (got {output_size})"
    assert isinstance(default_index, int), f"`default_index` should be an integer, (got {default_index})"
    assert isinstance(return_values, bool), f"`return_values` should be a boolean, (got {return_values})"

    dim = dim + input.dim() if dim < 0 else dim

    sizes = list(input.shape)
    if output_size is not None:
        sizes[dim] = output_size
    else:
        sizes[dim] = index.max().item() + 1 if index.numel() else 0

    if reduce == 'max':
        fill_val = torch.finfo(input.dtype).min
    else:
        fill_val = torch.finfo(input.dtype).max

    out_value = torch.full(sizes, fill_val, dtype=input.dtype, device=input.device)
    out_idx = torch.full(sizes, default_index, dtype=index.dtype, device=index.device)

    if input.numel() == 0:
        if return_values:
            return out_idx, out_value
        return out_idx

    offset1 = 1
    offset2 = 1
    for s in input.shape[:dim]:
        offset1 *= s
    for s in input.shape[dim + 1:]:
        offset2 *= s

    input = input.contiguous()
    index = index.contiguous()

    input_data = input.view(offset1, -1, offset2)
    index_data = index.view(offset1, -1, offset2)
    out_value_data = out_value.view(offset1, -1, offset2)
    out_idx_data = out_idx.view(offset1, -1, offset2)

    for i in range(offset1):
        for j in range(input.shape[dim]):
            for k in range(offset2):
                value = input_data[i, j, k]
                dim_index = index_data[i, j, k]
                assert 0 <= dim_index < out_idx.size(dim), f"Expected `input` values to be in the range {0} to " \
                        f"{out_idx.size(dim)}, (got {dim_index})"
                if reduce == 'max':
                    if out_value_data[i, dim_index, k] <= value:
                        out_value_data[i, dim_index, k] = value
                        out_idx_data[i, dim_index, k] = j
                else:
                    if out_value_data[i, dim_index, k] >= value:
                        out_value_data[i, dim_index, k] = value
                        out_idx_data[i, dim_index, k] = j

    if return_values:
        out_value.masked_fill_(out_value == fill_val, 0)
        return out_idx, out_value
    return out_idx
```

### Alternatives

_No response_

### Additional context

_No response_

cc @mikaylagawarecki