Create a PyTorch tensor of sequences which excludes specified value
I have a 1d PyTorch tensor containing integers between 0 and n-1. Now I need to create a 2d PyTorch tensor with n-1 columns, where each row is a sequence from 0 to n-1 excluding the value in the first tensor. How can I achieve this efficiently?
Ex:
n = 3
a = torch.Tensor([0, 1, 2, 1, 2, 0])
# desired output
b = [
    [1, 2],
    [0, 2],
    [0, 1],
    [0, 2],
    [0, 1],
    [1, 2]
  ]

Typically, the a.numel() &gt;&gt; n.
Detailed Explanation:
The first element of a is 0, hence it has to map to the sequence [0, 1, 2] excluding 0, which is [1, 2].
Similarly, the second element of a is 1, hence it has to map to [0, 2] and so on.
PS: I actually have an additional batch dimension, which I've excluded here for simplicity. Hence, I need the solution to be easily extendable to one additional dimension.

We can construct a tensor with the desired sequences and index with tensor a.
import torch

n = 3
a = torch.Tensor([0, 1, 2, 1, 2, 0]) # using torch.tensor is recommended

def exclude_gather(a, n):
    sequences = torch.nonzero(torch.arange(n) != torch.arange(n)[:,None], as_tuple=True)[1].reshape(-1, n-1)
    return sequences[a.long()]

exclude_gather(a, n)

Output
tensor([[1, 2],
        [0, 2],
        [0, 1],
        [0, 2],
        [0, 1],
        [1, 2]])

We can add a batch dimension with functorch.vmap
from functorch import vmap

n = 4
b = torch.Tensor([[0, 1, 2, 1, 3, 0],[0, 3, 1, 0, 2, 1]])

vmap(exclude_gather, in_dims=(0, None))(b, n)

Output
tensor([[[1, 2, 3],
         [0, 2, 3],
         [0, 1, 3],
         [0, 2, 3],
         [0, 1, 2],
         [1, 2, 3]],

        [[1, 2, 3],
         [0, 1, 2],
         [0, 2, 3],
         [1, 2, 3],
         [0, 1, 3],
         [0, 2, 3]]])

