Vectorizing double for loop

1

Is there a way to vectorize the “for” loops in the following scenario:

The variable “A” is a 2-dim tensor of shape (N, M).
The variable “B” is a 2-dim tensor of shape (M, M).

for k in range(N):

    for i in range(M):
        for j in range(M):
            B[i, j] = A[k, i] - A[k, j]
2

Based on the code it seems you are overwriting the intermediate results in B with the last value of k, so you could also remove the for k in range(N) loop.
To get rid of the loop you could unsqueeze the A tensor and use broadcasting.
Here is a code snippet showing the same results:

N, M  = 2, 3
A = torch.randn(N, M)
B = torch.randn(M, M)

for k in range(N):
    for i in range(M):
        for j in range(M):
            B[i, j] = A[k, i] - A[k, j]

C = torch.zeros(M, M)
for i in range(M):
    for j in range(M):
        C[i, j] = A[-1, i] - A[-1, j]

print(B - C)
# tensor([[0., 0., 0.],
#         [0., 0., 0.],
#         [0., 0., 0.]])

D = A[-1].unsqueeze(1) - A[-1].unsqueeze(0)

print(B - D)
# tensor([[0., 0., 0.],
#         [0., 0., 0.],
#         [0., 0., 0.]])
3

Thank you!
That helped.