PyTorch tensor assignment

1

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?

2

Direct indexing should work:

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)
1 Like