эквивалент питора tf.gather

Question

У меня возникли проблемы с переносом некоторого кода из тензорного потока в pytorch.

Итак, у меня есть матрица с размерами 10x30, представляющая 10 примеров, каждый с 30 функциями.Затем у меня есть другая матрица с размерами 10x5, содержащая индексы из 5 ближайших примеров для каждого примера в первой матрице.Я хочу «собрать», используя индексы, содержащиеся во второй матрице, 5 примеров шкафов для каждого примера в первой матрице, оставляя меня с трехмерным тензором формы 10x5x30.

В тензорном потоке это делается с помощью tf.gather(matrix1, matrix2),Кто-нибудь знает, как я мог сделать это в pytorch?

Answer 1

Как насчет этого?

matrix1 = torch.randn(10, 30)
matrix2 = torch.randint(high=10, size=(10, 5))
gathered = matrix1[matrix2]

Используется трюк индексации с массивом целых чисел.