Create a PyTorch tensor of sequences which excludes specified value

Question

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() >> 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.

Answer 1

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]]])

Answer 2

All you have to do is initialize a multi-dimension array with all possible indices using torch.arange(). After that, purge indices that you don't want from each tensor using a boolean mask.

import torch

a = torch.Tensor([0, 1, 2, 1, 2, 0])
n = 3
b = [torch.arange(n) for i in range(len(a))]
c = [b[i]!=a[i] for i in range(len(b))]
# use the boolean array as a mask to apply on b
d = [[b[i][c[i]] for i in range(len(b))]]

print(d) # this can be converted to a list of numbers or torch tensor

This prints the output - [[tensor([1, 2]), tensor([0, 2]), tensor([0, 1]), tensor([0, 2]), tensor([0, 1]), tensor([1, 2])]] which you can convert to int/numpy/torch array/tensor easily.

This can be extended to multiple dimensions as well.