PyTorch and NumPy implementation of pairwise distance function (p2dist) for tensors with dimensions greater or equal to 2. Distances are calculated w.r.t. last dimension.

p2dist.py

import torch
import numpy as np

def p2dist_pytorch(x, y):
    y_dim = len(y.shape)
    return torch.pow(x, 2).sum(dim=-1).view(x.shape[:-1]+(1,)) - \
        2 * torch.matmul(x, y.permute(list(range(y_dim-2))+[y_dim-1, y_dim-2])) + \
            torch.pow(y, 2).sum(dim=-1).view(y.shape[:-2]+(1,y.shape[-2]))

def p2dist_numpy(x, y):
    y_dim = len(y.shape)
    return np.power(x, 2).sum(axis=-1).reshape(x.shape[:-1]+(1,)) - \
        2 * np.matmul(x, np.transpose(y, (list(range(y_dim-2))+[y_dim-1, y_dim-2]))) + \
            np.power(y, 2).sum(axis=-1).reshape(y.shape[:-2]+(1,y.shape[-2]))

credits to @Fzaero