Select/Mask different column index in every row
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.

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

