Repository: pytorch/pytorch

Issue:
torch.unique is not work with `dim`, can not for batch unique 
### 🐛 Describe the bug

```
import torch
t = torch.randn(2, 3, 4)
t[t>0.2] = 1
t[t<=0.2] = 0
print(torch.unique(t))
print(torch.unique(t, dim=-1))
```

the expected output shape for `print(torch.unique(t, dim=-1))` should be `[2, 3, 2]`, but it output the same shape with original tensor.

### Versions

PyTorch version: 2.0.1
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.10.10 (main, Mar 21 2023, 18:45:11) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.13.0-30-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA GeForce RTX 3090
GPU 1: NVIDIA GeForce RTX 3090

Nvidia driver version: 520.61.05
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Byte Order:                      Little Endian
Address sizes:                   46 bits physical, 48 bits virtual
CPU(s):                          36
On-line CPU(s) list:             0-35
Thread(s) per core:              2
Core(s) per socket:              18
Socket(s):                       1
NUMA node(s):                    1
Vendor ID:                       GenuineIntel
CPU family:                      6
Model:                           85
Model name:                      Intel(R) Core(TM) i9-10980XE CPU @ 3.00GHz
Stepping:                        7
CPU MHz:                         3000.000
CPU max MHz:                     4800.0000
CPU min MHz:                     1200.0000
BogoMIPS:                        6000.00
Virtualization:                  VT-x
L1d cache:                       576 KiB
L1i cache:                       576 KiB
L2 cache:                        18 MiB
L3 cache:                        24.8 MiB
NUMA node0 CPU(s):               0-35
Vulnerability Itlb multihit:     KVM: Mitigation: VMX disabled
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Enhanced IBRS, IBPB conditional, RSB filling
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Mitigation; TSX disabled
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl vmx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb cat_l3 cdp_l3 invpcid_single ssbd mba ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid cqm mpx rdt_a avx512f avx512dq rdseed adx smap clflushopt clwb intel_pt avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local dtherm ida arat pln pts hwp hwp_act_window hwp_epp hwp_pkg_req avx512_vnni md_clear flush_l1d arch_capabilities

Versions of relevant libraries:
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.23.5
[pip3] pytorch-metric-learning==2.0.1
[pip3] torch==2.0.1
[pip3] torchaudio==2.0.2
[pip3] torchvision==0.15.2
[pip3] triton==2.0.0
[conda] blas                      1.0                         mkl  
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-service               2.4.0           py310h7f8727e_0  
[conda] mkl_fft                   1.3.1           py310hd6ae3a3_0  
[conda] mkl_random                1.2.2           py310h00e6091_0  
[conda] numpy                     1.23.5          py310hd5efca6_0  
[conda] numpy-base                1.23.5          py310h8e6c178_0  
[conda] pytorch                   2.0.1           py3.10_cuda11.8_cudnn8.7.0_0    pytorch
[conda] pytorch-cuda              11.8                 h7e8668a_5    pytorch
[conda] pytorch-metric-learning   2.0.1              pyhd8ed1ab_0    conda-forge
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchaudio                2.0.2               py310_cu118    pytorch
[conda] torchtriton               2.0.0                     py310    pytorch
[conda] torchvision               0.15.2              py310_cu118    pytorch

cc @ezyang @gchanan @zou3519 @albanD @svekars @carljparker


Discussion:
From quick check, it seems that no value for dim= returns the expected shape right both on CPU and CUDA.
Raising priority to investigate.
Keeping high pri to investigate, but for consistency sake we want to preserve Python behavior of treating negative dims as counted from the end of the tensor dims
I think it is actually working properly, and I think that the assertion that "the expected output shape for `print(torch.unique(t, dim=-1))` should be `[2, 3, 2]`" is not correct. The output shape will depend on the values--even if all elements are only either 1 or 0.

The Pytorch documentation does not describe the behavior of the `dim` arg well. It only says "dim (int) - the dimension to apply unique. If None, the unique of the flattened input is returned", which is quite ambiguous for the `int` case.

But Numpy's documentation describes it in pretty good detail, particularly in the "Notes" section: https://numpy.org/doc/stable/reference/generated/numpy.unique.html

According to that explanation, given an input array of shape `(2, 3, 4)` (like the one given in the issue description) and `dim=-1`, `unique` will first create a temporary array where the last dimension is permuted to the first dimension, giving shape `(4, 2, 3)`. Then all the dimensions except for the first are temporarily flattened, giving shape `(4, 6)`.  Now we have 4 individual vectors of size 6. The vectors (not the elements within the vectors) will be sorted, and then any duplicate vectors are thrown out, leaving us with either size `(1, 6)`, `(2, 6)`, `(3, 6)`, or `(4, 6)`, depending on how many duplicates there were. Then the original shapes of the remaining vectors will be restored from `(6,)` to `(2, 3)`. Finally the first dimension will be moved back to the last dimension.

The possible shapes of the final result are: `(2, 3, 4)` if the 4 vectors were all unique, `(2, 3, 3)` if there was one duplicate, `(2, 3, 2)` if there were two duplicates, or `(2, 3, 1)` if there were three duplicates.

Here's a specific example, like the one in the issue description:
```python
>>> import torch
>>> torch.manual_seed(123)
>>> t = torch.rand(2, 3, 4)
>>> t[t>0.2] = 1
>>> t[t<=0.2] = 0
>>> t
tensor([[[1., 1., 1., 1.],
         [0., 1., 0., 0.],
         [0., 1., 1., 1.]],

        [[0., 0., 1., 1.],
         [0., 1., 1., 1.],
         [1., 1., 1., 1.]]])
```

Here, if we give `dim=-1`, then we'll have the 4 6-element vectors that we can see by drawing four vertical lines through the text printout of the array, namely `[1, 0, 0, 0, 0, 1]`, `[1, 1, 1, 0, 1, 1]`, `[1, 0, 1, 1, 1, 1]`, and `[1, 0, 1, 1, 1, 1]`. Those last two arrays are identical, so one of them will be removed, and the final result will have shape `(2, 3, 3)`. And that's what we get:

```python
>>> torch.unique(t, dim=-1)
tensor([[[1., 1., 1.],
         [0., 0., 1.],
         [0., 1., 1.]],

        [[0., 1., 0.],
         [0., 1., 1.],
         [1., 1., 1.]]])
>>> torch.unique(t.cuda(), dim=-1)
tensor([[[1., 1., 1.],
         [0., 0., 1.],
         [0., 1., 1.]],

        [[0., 1., 0.],
         [0., 1., 1.],
         [1., 1., 1.]]], device='cuda:0')
>>> import numpy as np
>>> np.unique(t, axis=-1)
array([[[1., 1., 1.],
        [0., 0., 1.],
        [0., 1., 1.]],

       [[0., 1., 0.],
        [0., 1., 1.],
        [1., 1., 1.]]], dtype=float32)
```

I think probably we should improve the documentation to better describe this behavior
Also stumbled on this. The docs don't have any example of this :(

But it would be great to have unique supporting batch mode (maybe returning results in some sort of padded or masked format - e.g. `return_inverse = True` by itself can certainly be done in batch mode, and unique items by themselves can be returned in some padded/NestedTensor format)
Let's update the doc with more details similar to the numpy one then!

Resolving patch:
diff --git a/torch/functional.py b/torch/functional.py
--- a/torch/functional.py
+++ b/torch/functional.py
@@ -760,8 +760,11 @@ def _unique_impl(input: Tensor, sorted: bool = True,
             elements in the original input ended up in the returned unique list.
         return_counts (bool): Whether to also return the counts for each unique
             element.
-        dim (int): the dimension to apply unique. If ``None``, the unique of the
-            flattened input is returned. default: ``None``
+        dim (int, optional): the dimension to operate upon. If ``None``, the
+            unique of the flattened input is returned. Otherwise, each of the
+            tensors indexed by the given dimension is treated as one of the
+            elements to apply the unique operation upon. See examples for more
+            details. Default: ``None``
 
     Returns:
         (Tensor, Tensor (optional), Tensor (optional)): A tensor or a tuple of tensors containing
@@ -799,6 +802,76 @@ def _unique_impl(input: Tensor, sorted: bool = True,
         tensor([[0, 2],
                 [1, 2]])
 
+        >>> a = torch.tensor([
+        ...     [
+        ...         [1, 1, 0, 0],
+        ...         [1, 1, 0, 0],
+        ...         [0, 0, 1, 1],
+        ...     ],
+        ...     [
+        ...         [0, 0, 1, 1],
+        ...         [0, 0, 1, 1],
+        ...         [1, 1, 1, 1],
+        ...     ],
+        ...     [
+        ...         [1, 1, 0, 0],
+        ...         [1, 1, 0, 0],
+        ...         [0, 0, 1, 1],
+        ...     ],
+        ... ])
+
+        >>> # If we call `torch.unique(a, dim=0)`, each of the tensors `a[idx, :, :]`
+        >>> # will be compared. We can see that `a[0, :, :]` and `a[2, :, :]` match
+        >>> # each other, so one of them will be removed.
+        >>> (a[0, :, :] == a[2, :, :]).all()
+        tensor(True)
+        >>> a_unique_dim0 = torch.unique(a, dim=0)
+        >>> a_unique_dim0
+        tensor([[[0, 0, 1, 1],
+                 [0, 0, 1, 1],
+                 [1, 1, 1, 1]],
+                [[1, 1, 0, 0],
+                 [1, 1, 0, 0],
+                 [0, 0, 1, 1]]])
+
+        >>> # Notice which sub-tensors from `a` match with the sub-tensors from
+        >>> # `a_unique_dim0`:
+        >>> (a_unique_dim0[0, :, :] == a[1, :, :]).all()
+        tensor(True)
+        >>> (a_unique_dim0[1, :, :] == a[0, :, :]).all()
+        tensor(True)
+
+        >>> # For `torch.unique(a, dim=1)`, each of the tensors `a[:, idx, :]` are
+        >>> # compared. `a[:, 0, :]` and `a[:, 1, :]` match each other, so one of
+        >>> # them will be removed.
+        >>> (a[:, 0, :] == a[:, 1, :]).all()
+        tensor(True)
+        >>> torch.unique(a, dim=1)
+        tensor([[[0, 0, 1, 1],
+                 [1, 1, 0, 0]],
+                [[1, 1, 1, 1],
+                 [0, 0, 1, 1]],
+                [[0, 0, 1, 1],
+                 [1, 1, 0, 0]]])
+
+        >>> # For `torch.unique(a, dim=2)`, the tensors `a[:, :, idx]` are compared.
+        >>> # `a[:, :, 0]` and `a[:, :, 1]` match each other. Also, `a[:, :, 2]` and
+        >>> # `a[:, :, 3]` match each other as well. So in this case, two of the
+        >>> # sub-tensors will be removed.
+        >>> (a[:, :, 0] == a[:, :, 1]).all()
+        tensor(True)
+        >>> (a[:, :, 2] == a[:, :, 3]).all()
+        tensor(True)
+        >>> torch.unique(a, dim=2)
+        tensor([[[0, 1],
+                 [0, 1],
+                 [1, 0]],
+                [[1, 0],
+                 [1, 0],
+                 [1, 1]],
+                [[0, 1],
+                 [0, 1],
+                 [1, 0]]])
     """
     if has_torch_function_unary(input):
         return handle_torch_function(


Test patch: