Сумма по показателям со значением (1)

Question

У меня есть тензор с формой (2, 2, 3), например:

a= tensor ([[[2, 0, 2],[1, 0, 0]],[[1, 0, 1],[0, 1, 0]]])

Я хочу найти индексы значений (1), затем сделать 2 в степени этих индексов и, наконец, добавить результаты для последнего измерения, поэтому конечный результат должен выглядеть следующим образом:

tensor ([[[],[2^0]],[[2^0+2^2],[2^1]]])

Мой фактический тензор намного больше, чем в этом примере, поэтому я не хочу использовать «для» l oop, и Я должен использовать трансляцию ...

Я думал о чем-то вроде torch.pow(2,(a == 1).nonzero()).sum(), но это не работает. Мне нужно найти способ применить (a== 1).nonzero() только для последнего измерения, есть предложения? Спасибо.

Answer 1

Измените его на NumPy и примените это:

a_n = a.numpy()
a_n = np.apply_along_axis(func1d=lambda x: np.sum(np.power(2,np.where(x==1))[0]), axis=2, arr=a_n)

a = torch.Tensor(a_n)

В основном, это применяет функцию, которую вы хотите в оси = 2, я предполагаю, что ваш больший массив что-то вроде (2,2, п)