Select/Mask different column index in every row

user_1994382 0
Viewed 6487

In pytorch I have a multi-dimensional tensor, call it X

X = [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12], ...]

Now I would like to select a different column index for each row like so

indices = [[0], [1], [0], [2], ...]
# now I expect following values to be returned:
[[1], [5], [7], [12], ...]

also I would like to achieve the opposite so that for the given indices I get

[[2, 3], [4, 6], [8, 9], [10, 11]]

Is there a "simple" way to achieve this without a for loop? I would be grateful for any ideas.

pytorch
2 Answers

In fact the torch.gather function performs exactly this.

For example

a = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])
indices = torch.tensor([[0], [1], [0], [2]])
a.gather(1, indices)

will exactly return

tensor([[ 1],
        [ 5],
        [ 7],
        [12]])

I do not require the opposite anymore but for this I would propose just creating a mask with all ones and then setting the respective indices of the "gathering" tensor to 0 or just create a new "gathering" tensor which contains the respective opposite keys. For example:

indices_opposite = [np.setdiff1d(np.arange(a.size(1)), i) for i in indices.numpy()]
user_19943820

It is easily done with numpy, it can convert forth and back if you need it for tensor 

X = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])
X2= X[np.arange(X.shape[0]),[0,1,0,2]]
X3 = np.setdiff1d(X, X2).reshape(X.shape[0],X.shape[1]-1)

EDIT

doing the same in gpu, working with tensors

import torch.tensor as tensor
import numpy as np

X = tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])
X2= X[np.arange(X.shape[0]),[0,1,0,2]]

def allothers(X, num):
    return [x for x in range(X.shape[1]) if x not in [num]]


X3 = X[ [[x] for x in np.arange(X.shape[0])], [allothers(X, 0),allothers(X, 2),allothers(X, 1),allothers(X, 1)] ]
user_85636490
Related