In PyTorch calc Euclidean distance instead of matrix multiplication

Question

Let say we have 2 matrices:

mat = torch.randn([20, 7]) * 100
mat2 = torch.randn([7, 20]) * 100

n, m = mat.shape

The simplest usual matrix multiplication looks like this:

def mat_vec_dot_product(mat, vect):
    n, m = mat.shape
    
    res = torch.zeros([n])
    for i in range(n):
        for j in range(m):
            res[i] += mat[i][j] * vect[j]
        
    return res

res = torch.zeros([n, n])
for k in range(n):
    res[:, k] = mat_vec_dot_product(mat, mat2[:, k])
    

But what if I need to apply L2 norm instead of dot product? The code is next:

def mat_vec_l2_mult(mat, vect):
    n, m = mat.shape
    
    res = torch.zeros([n])
    for i in range(n):
        for j in range(m):
            res[i] += (mat[i][j] - vect[j]) ** 2
            
    res = res.sqrt()
        
    return res

for k in range(n):
    res[:, k] = mat_vec_l2_mult(mat, mat2[:, k])

Can we do this somehow in an optimal way using Torch or any other libraries? Cause naive O(n^3) Python code works really slow.

Answer 1

Use torch.cdist for L2 norm - euclidean distance

res = torch.cdist(mat, mat2.permute(1,0), p=2)

Here, I have used permute to swap dim of mat2 from 7,20 to 20,7