A bug in torch.unique with return_inverse=True and dim != None?

1

Try this:

aa = torch.tensor([[1, 2, 2], [3, 3, 3]])
unq, groups = aa.unique_consecutive(dim=1, return_inverse=True)

groups is:
tensor([0, 1, 1])

which are the group assignments for the first row only! expected groups to be of shape [2, 3], with the group assignments for each row… Is this a bug? Any way of achieving what I intended in an efficient way?

2

Hi,

You are reducing your function over second dimension, so it has to be in shape of [3, ] in your case.

Also, it was just a coincidence that the returned groups is same as groups for just the first row!
When you set the dim=1, your function will be applied on given dimension rather than whole tensor as a flattened array. So, set dim=None.

You can run the following code in 4 cases,

  • with aa=aa1 and dim=None
  • With aa=aa1 and dim=1
  • with aa=aa2 and dim=None
  • With aa=aa2 and dim=1
aa1 = torch.tensor([[1, 0, 0],
                   [0, 0, 1]])
aa2 = torch.tensor([[1, 1, 1],
                   [0, 1, 1]])
aa = aa1  # change this to aa2 with dim=1 and dim=None and see the difference
print(aa)
unq, groups = aa.unique_consecutive(dim=1, return_inverse=True)
print(groups)
unq, groups = aa.unique_consecutive(dim=None, return_inverse=True)
print(groups)

Bests