Find euclidean distance between a tensor and each row tensor of a matrix efficiently in PyTorch

dsplnme 0
Viewed 50

I have a tensor A of size torch.Size([3]) and another tensor B of size torch.Size([4,3]).

I want to find the distance between A and each of the 4 rows of B.

I'm new to Torch and I reckon a for loop for each of the rows wouldn't be efficient. I have looked into torch.linalg.norm and torch.cdist but I'm not sure if they solve my problem, unless I'm missing something.

euclidean-distance torch pytorch python
1 Answers

You look for:

torch.norm(A[None, :] - B, p=2, dim=1)
  • A[None, :] resize the tensor to shape (1, 3)
  • A[None, :] - B will copy 4 times the tensor A to match the size of B ("broadcast") and make the substraction
  • torch.norm(..., p=2, dim=1) computes the euclidian norme for each column.

Output shape: (4,)

Valentin Goldité0
Related