Как выбрать субтенор из тензора

Question

У меня есть тензор A с размерностью [N1, N2, N3/2, 2, N4, N5]. Другой тензор B - это тензор булева индекса с размерностью [N1, N2, N3/2] = 1,0,0,0,1....
Теперь я хочу получить субтенор C с размером [N1, N2, N3/2, N4, N5], использующий индекс B.
Как я могу реализовать это с Pytorch?
Обратите внимание, что я не склонен использовать цикл for, потому что он очень медленный.

Я просмотрел некоторые функции в pytorch, но обнаружил, что ни одна из них не подходит для моей задачи

Answer 1

Вы не имеете смысла: Как C может иметь те же shape, что и A, когда B выбирает только часть записей A для копирования в C ??

Предположим, что B имеет k ненулевых элементов (то есть B выбирает k элементов из N1*N2*N3/2 матриц размера N4*N5 в A), тогда C может иметь форму [k, N4, N5] только с k < N1*N2*N3/2.

Вы можете использовать nonzero() для преобразования B из логических индексов в k целочисленные индексы, а затем использовать их для выбора k элементов

C = A.view(-1, *A.shape[-2:])[B.nonzero(), ...]