Sum over indices with value (1)

user_13210809 0
Viewed 92

I have a tensor with a shape (2, 2, 3) like:

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

I want to find the indices of the values (1), then make 2 to the power of those indices and finally add the results for the last dimension, so the final result should be like:

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

My actual tensor is much bigger than this example, so I don’t want to use "for" loop, and I have to use broadcasting…

I was thinking of something like torch.pow(2,(a == 1).nonzero()).sum(), but it doesn’t work. I have to find a way to apply (a== 1).nonzero() only for the last dimension, any suggestion? Thanks.

pytorch python
1 Answers

Change it to NumPy and apply this:

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)

Basically, it applies the function you want in axis=2, I'm supposing your larger array is something like (2,2,n)

user_102602430
Related