-
Notifications
You must be signed in to change notification settings - Fork 26.3k
Description
🚀 Feature
The current implementation of torch.unique does not have a return_counts argument, which returns the counts of all unique occurrences in the tensor.
Motivation
The feature request is related to some users missing this option, e.g. this thread. A current workaround is to use something like this:
x = torch.tensor([1, 2, 2, 2, 3])
x_unique = x.unique(sorted=True)
x_unique_count = torch.stack([(x==x_u).sum() for x_u in x_unique])
Also, numpy supports this option, which should be a nice to have for new users. np.unique
As this might be added as the last argument, no breaking changes would occur.
Pitch
I could add this option to the current torch.unique implementation as I've recently added the return_inverse option and am still familiar with the code.
Alternatives
An alternative would be to use the hacky code from above.
Let me know, if that sound good to you!
Also a bit unrelated, but we could also add the return_index option together to match the numpy function as close as possible, but that's probably another feature request.
Best,
ptrblck