Add return_counts to torch.unique (#18391)
authorXiang Gao <qasdfgtyuiop@gmail.com>
Tue, 26 Mar 2019 03:30:33 +0000 (20:30 -0700)
committerFacebook Github Bot <facebook-github-bot@users.noreply.github.com>
Tue, 26 Mar 2019 03:38:17 +0000 (20:38 -0700)
commite2730ddb21164f1cb36330289ab2cd0629c06dc8
treeaddaf20e6a382ccee4c1c599264bf2e9cb5b7e3e
parentba4de667fa578c84a255da3ef57948e14af7f01a
Add return_counts to torch.unique (#18391)

Summary:
Fixes: https://github.com/pytorch/pytorch/issues/12598

This PR was originally authorized by ptrblck at https://github.com/pytorch/pytorch/pull/15495, but since there was no update for months after the request change, I clone that branch and resolve the code reviews here. Hope everything is good now. Especially, the implementation of count is changed from ptrblck's original algorithm to the one ngimel suggest, i.e. using `unique_by_key` and `adjacent_difference`.

The currently implementation of `_unique_dim` is VERY slow for computing inverse index and counts, see https://github.com/pytorch/pytorch/issues/18405. I will refactor `_unique_dim` in a later PR. For this PR, please allow me to keep the implementation as is.

cc: ptrblck ezyang ngimel colesbury
Pull Request resolved: https://github.com/pytorch/pytorch/pull/18391

Reviewed By: soumith

Differential Revision: D14605905

Pulled By: VitalyFedyunin

fbshipit-source-id: 555f5a12a8e28c38b10dfccf1b6bb16c030bfdce
aten/src/ATen/native/Unique.cpp
aten/src/ATen/native/cuda/Unique.cu
aten/src/ATen/native/native_functions.yaml
aten/src/ATen/native/sparse/SparseTensor.cpp
test/test_torch.py
tools/autograd/derivatives.yaml
torch/functional.py
torch/onnx/symbolic.py
torch/tensor.py