I'm doing this question where I am given two PyTorch matrices A and B of shapes (M,d) and (N,d) and I want to create a matrix of the distances between each pair of rows. So the resulting distance matrix has D[i,j] = torch.sum((A[i] - B[j])**2)
I want to write the code in a vectorized manner. The suggested hint was to use a matrix multiplication, and that made me think of the solution that worked: since AB^T would give the dot product of row i of A with row j of B, I could use the fact that |v - w|^2 = |v|^2 + |w|^2 - 2|v.w| to get a solution that works.
dists = (torch.sum(torch.square(A),dim=1).view((-1,1))
+ torch.sum(torch.square(B),dim=1).view((1,-1))
- 2*A @ B.t())
However, this was not my initial solution and my initial solution, which also seems vectorized, runs about 100 times slower than the above code (both solutions have been tested and are correct).
My initial solution involved broadcasting. Since I want to do |A[i] - B[j]|^2, I just reshaped A to have shape (M,d,1) and B to have shape (1,d,N) so that by broadcasting, I would be calculating all the possible A[i]-B[j]s. Then it was just a matter of squaring and summing about the middle axis to go from shape (M,d,N) -> (M,N)
A_v = A.view((A.shape[0],-1,1))
B_v = B.view((B.shape[0],-1)).permute((1,0))
dists=torch.sum(torch.square(A_v-B_v),dim=1)
Can you help me understand why this code is not vectorized and doesn't run as efficiently on the GPU? If possible, give me a general idea of how to know in future that code like this is not really vectorized in the same way as the one above.