To calculate euclidean distance between vectors in a torch tensor with multiple dimensions

Question

There is a random initialized torch tensor of the shape as below.

Inputs

tensor1 = torch.rand((4,2,3,100))
tensor2 = torch.rand((4,2,3,100))

tensor1 and tensor2 are torch tensors with 24 100-dimensional vectors, respectively.

I want to get a tensor with a shape of torch.size([4,2,3]) by obtaining the Euclidean distance between vectors with the same index of two tensors.

I used dist = torch.nn.functional.pairwise_distance(tensor1, tensor2) to get the results I wanted.

However, the pairwise_distance function calculates the euclidean distance for the second dimension of the tensor. So dist shape is torch.size([4,3,100]).

I have performed transpose several times to solve these problems. My code is as follows.

tensor1 = tensor1.transpose(1,3)
tensor2 = tensor2.transpose(1,3)
dist = torch.nn.functional.pairwise_distance(tensor1, tensor2)
dist = dist.transpose(1,2)

Is there a simpler or easier way to get the result I want?

Answer 1

Here ya go

dist = (tensor1 - tensor2).pow(2).sum(3).sqrt()

Basically that's what Euclidean distance is.

Subtract -> power by 2 -> sum along the unfortunate axis you want to eliminate-> square root