Transforming the indices back through view in pytorch
I have the following code, where I define a torch tensor of zeros, change one item to be equal to 1, pass through three reshape functions.
Then, after all the transformations, I obtain the index of 1. I am wondering if it is possible to somehow use the max_idx and the information about the permutations/.view to obtain the index of 1 in the initial B1 tensor (which should be equal to 1234).
A1 = np.zeros(10*18*40*28)
A1[1234] = 1
A1 = A1.reshape(10, 18, 40, 28)
B1 = torch.Tensor(A1)
print('B1: ', B1.shape, torch.nonzero(B1))
C1 = B1.permute(0, 2, 3, 1)
print('C1: ', C1.shape,  torch.nonzero(C1))
D1 = C1.contiguous().view(C1.shape[0], C1.shape[1], C1.shape[2], 3, 6)
print('D1: ', D1.shape, torch.nonzero(D1))
E1 = D1.contiguous().view(D1.shape[0], -1, 6)
print('E1: ', E1.shape, torch.nonzero(E1))

max_idx = torch.nonzero(E1)

I would love to hear any hints on how I can try to do this :)

For each dimension, you can check in which index the value you refer to lies. After you find the right index, you subtract the number of indices that came before the next dimension and solve the same problem for a smaller sub-array.
Or you can just use the function 'numpy.unravel_index' that does the exact same thing.
import numpy as np
import torch


A1 = np.zeros(10*18*40*28)
idx = 1234
A1[idx] = 1
A1 = A1.reshape(10, 18, 40, 28)
B1 = torch.Tensor(A1)
print('B1: ', B1.shape, torch.nonzero(B1))

idx_temp = idx+0
idxB1 = np.zeros((B1.dim(),), dtype = int)
for i in range(B1.dim()):
    idxB1[i]  = idx_temp//np.prod(B1.shape[i+1:])
    idx_temp -= np.prod(B1.shape[i+1:])*idxB1[i]

idxB1np = np.unravel_index(idx, B1.shape)

print(f'idxB1 = {idxB1}')
print(f'idxB1np = {idxB1np}')

output:
B1:  torch.Size([10, 18, 40, 28]) tensor([[0, 1, 4, 2]])
idxB1 = [0 1 4 2]
idxB1np = (0, 1, 4, 2)

