Выберите / замаскируйте разные индексы столбцов в каждой строке

Question

В pytorch у меня есть многомерный тензор, назовите его X

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

Теперь я хотел бы выбрать другой индекс столбца для каждой строки, например,

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

Я хотел бы добиться обратного, чтобы для данных индексов я получил

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

Есть ли "простой" способ добиться этого без цикла for? Буду благодарен за любые идеи.

Answer 1

Это легко сделать с помощью NumPy, он может конвертировать туда и обратно, если вам это нужно для тензора

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