pytorch best way to call unique_consecutive() on certain dimension of a tensor?

Question

What's the best way to call unique_consecutive() on certain dimension of a tensor, and pad the left-out with specified values? For simplicity, we could use 2D tensor as an example:

input = tensor([[3, 3, 5, 5, 5],
                [3, 3, 2, 2, 3]])

if specifying padding value -1, what i hope to get is:

output = tensor([[3, 5, -1, -1, -1],
                 [3, 2, 3, -1, -1]])

Also, to save memory, i'd prefer to do so in-place if possible. Thanks!

user_276625 · Answer

found a decent solution from PyTorch forum https://discuss.pytorch.org/t/best-way-to-run-unique-consecutive-on-certain-dimension/112662, quoted here for quick reference.


    x = tensor([[3, 3, 5, 5, 5],
                [3, 3, 2, 2, 3]])
    unique_x, indices = torch.unique_consecutive(x, return_inverse=True)
    indices -= indices.min(dim=1, keepdims=True)[0]
    result = -torch.ones_like(x)
    result = result.scatter_(1, indices, x)

pytorch best way to call unique_consecutive() on certain dimension of a tensor?

1 Answers