Hi everyone, I’m trying to assign elements from a 6D tensor x, to a 4D tensor y. I would like to optimize this process. I write here the code with for cycles:
for r in range(R):
for k in range(K):
y[r,k,:,:] = x[r,r,k,k,:,:]
the dimensions of tensor are:
x.size()=(R,R,K,K, N, M)
y.size()=(R,K,N,M)
How can I do this avoiding explicit for cycles?
R, K, N, M = 2, 3, 4, 5
x = torch.randn(R, R, K, K, N, M)
y = torch.zeros(R, K, N, M)
# reference
for r in range(R):
for k in range(K):
y[r,k,:,:] = x[r,r,k,k,:,:]
# direct indexing
z = x[torch.arange(R)[:, None], torch.arange(R)[:, None], torch.arange(K), torch.arange(K)]
# comapre
print((z == y).all())
# tensor(True)