Repository: pytorch
Issue: Random Shuffle along Axis
Labels: feature,triaged,module: random

### 🚀 The feature, motivation and pitch

Dear PyTorch Devs,

Thank you for your hard work and dedication to creating a great ecosystem of tools and community of users.

This feature request proposes adding a standard lib function to shuffle "rows" across an axis of a tensor.  We will consider two ways to do this (use-case 1 and use-case 2).   Consider that an axis partitions a tensor into sub-tensors.  For instance a tensor of shape (2,3,4)  has 3 sub-tensors at axis=1.

Example input:
```
>>>  x = torch.arange(2*3*4).reshape(2,3,4)
>>>  x
tensor([[[ 0,  1,  2,  3],
         [ 4,  5,  6,  7],
         [ 8,  9, 10, 11]],

        [[12, 13, 14, 15],
         [16, 17, 18, 19],
         [20, 21, 22, 23]]])
```

Use-case 1:  Randomly shuffle an axis the same way for all sub-tensors.  (This implementation is trivially "solved" below).

```
>>> axis = 1  # choose an axis
>>> x.index_select(axis, T.randperm(x.shape[axis]))
tensor([[[ 0,  1,  2,  3],
         [ 8,  9, 10, 11],
         [ 4,  5,  6,  7]],

        [[12, 13, 14, 15],
         [20, 21, 22, 23],
         [16, 17, 18, 19]]])
```

Use-case 2: Randomly shuffle an axis differently for each other axes
```
>>> shufflerow(x, axis=1)
tensor([[[ 0,  1,  2,  3],
         [ 4,  5,  6,  7],
         [ 8,  9, 10, 11]],

        [[16, 17, 18, 19],
         [12, 13, 14, 15],
         [20, 21, 22, 23]]])
```

Where the `shufflerow` function (for use-case 2) could be implemented like this
```
def shufflerow(tensor, axis):
    row_perm = torch.rand(tensor.shape[:axis+1]).argsort(axis)  # get permutation indices
    for _ in range(tensor.ndim-axis-1): row_perm.unsqueeze_(-1)
    row_perm = row_perm.repeat(*[1 for _ in range(axis+1)], *(tensor.shape[axis+1:]))  # reformat this for the gather operation
    return tensor.gather(axis, row_perm)
```

- A more detailed explanation of the solution of use-case 2 is here on Stack Overflow if above is confusing:  https://stackoverflow.com/a/70758131/824039


It would be convenient to have a simple function like `shuffle(tensor, axis, mode:str)` for both use cases, where for example, we might say `mode='same'`  (use case 1) or `mode='different'` (use case 2).

### Alternatives

I do not know of a better shuffle implementation for use-case 2, but my proposed implementation seems not efficient.  Maybe an improvement would have an implementation in C or cuda that covers use-case 2 without the for loops and without the repeat and gather.

Currently, we have **torch.randperm** to randomly shuffle one axis the same way across all the same way.  

Perhaps off topic comment: I also wish PyTorch (and NumPy) had a toolkit dedicated to sampling, such as reservoir sampling across minibatches.  Sampling often introduces subtle bugs.

### Additional context

Variations of this feature request on the internet:

When axis=-1:
- https://discuss.pytorch.org/t/what-is-the-most-efficient-way-to-shuffle-each-row-of-a-tensor-with-different-shuffling-order-for-each-of-the-row/109772

For 2-d tensor:  
- https://stackoverflow.com/questions/44738273/torch-how-to-shuffle-a-tensor-by-its-rows

ChannelShuffle:
- https://pytorch.org/docs/stable/_modules/torch/nn/modules/channelshuffle.html#ChannelShuffle

Thank you!

cc @pbelevich