import torch


def squared_dist(x, y):
  """Calculate squared Euclidean distance
  Args:
    x: tensor of size (N, D) 
    y: tensor of size (M, D)
    
  Returns:
    d: tensor of size (N, M)
    
  Examples:
    x = torch.randn(3, 7)
    y = torch.randn(5, 7)
    d = dist(x, y)
    torch.norm(d - torch.norm(x.unsqueeze(1)-y, dim=-1)**2)
  """
  return torch.sum(x*x, dim=-1, keepdim=True) + torch.sum(y*y, dim=-1) - 2 * torch.mm(x, y.t())