Repository: pytorch
Issue: torch.unique(x, dim=1, return_inverse=True) returns inverse for only the last sub-tensor along dim
Labels: module: docs,triaged,module: numpy,module: sorting and selection

## Issue description

`torch.unique` called with the dim argument and return_inverse=True returns inverse for only the last sub-tensor along dimension dim. Also the first return value is not unique, so it seems the expected behavior should be documented (or there is a bug to be fixed).

## Code example

```
>>> x = torch.tensor([[1,1,2,3,4,2,2],[1,1,2,2,3,1,4]])
>>> uniq, inv = torch.unique(x, return_inverses=True, dim=1)
#Returned values are not unique:
>>> uniq
tensor([[1, 2, 2, 2, 3, 4],
            [1, 1, 2, 4, 2, 3]]
#inverse is returned only for the last index along dim:
>>> inv
tensor([0, 0, 2, 4, 5, 1, 3])
>>> (uniq[1,inv] == x[1]).all()
tensor(True)
```

I don't understand what `uniq` is supposed to be in this case. Clearly, it doesn't hold unique values. However, the behavior in torch seems to match that of numpy:
```
>>> np.unique(x.numpy(), return_inverse=True, axis=1)
(array([[1, 2, 2, 2, 3, 4],
       [1, 1, 2, 4, 2, 3]]), array([0, 0, 2, 4, 5, 1, 3]))
```

The behavior I expected would be something similar to looping over the outer dimension, i.e:
```
>>> [torch.unique(x[i], return_inverse=True) for i in range(2)]
[(tensor([1, 2, 3, 4]), tensor([0, 0, 1, 2, 3, 1, 1])), 
 (tensor([1, 2, 3, 4]), tensor([0, 0, 1, 1, 2, 0, 3]))]

```

Now, clearly there's a problem with returning unique values for each dimension in a single tensor, since the number of uniques may differ for each index. Maybe a single rectangular tensor with filler values towards the end. Or a list of 1d tensors. 

I don't know what the first return value is supposed to be in the current behavior, but since it matches with numpy I figure perhaps there's some useful meaning to it. I think it should be documented what the current behavior is, or it should be changed to match the loop semantics and the inverses should be returned for all indexes.


## System Info
PyTorch version: 1.8.0a0+d4c1684
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.10 (x86_64)
GCC version: (Ubuntu 10.2.0-13ubuntu1) 10.2.0
Clang version: Could not collect
CMake version: version 3.18.2

Python version: 3.8 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: GeForce RTX 3090
Nvidia driver version: 455.38
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.18.5
[pip3] numpydoc==1.1.0
[pip3] torch==1.8.0a0
[pip3] torchaudio==0.8.0a0+f1d3e1a
[pip3] torchvision==0.9.0.dev20201105
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.0.221             h6bb024c_0  
[conda] magma-cuda110             2.5.2                         1    pytorch
[conda] mkl                       2020.1                      217  
[conda] mkl-include               2020.2                      256  
[conda] mkl-service               2.3.0            py38he904b0f_0  
[conda] mkl_fft                   1.1.0            py38h23d657b_0  
[conda] mkl_random                1.1.1            py38h0573a6f_0  
[conda] numpy                     1.18.5           py38ha1c710e_0  
[conda] numpy-base                1.18.5           py38hde5b4d6_0  
[conda] numpydoc                  1.1.0                      py_0  
[conda] pytorch                   1.8.0.dev20201105 py3.8_cuda11.0.221_cudnn8.0.3_0    pytorch-nightly
[conda] torch                     1.8.0a0                  pypi_0    pypi
[conda] torchaudio                0.8.0.dev20201105            py38    pytorch-nightly
[conda] torchvision               0.9.0.dev20201105      py38_cu110    pytorch-nightly



cc @jlin27 @mruberry @rgommers @heitorschueroff