import torch


def gaussian_kernel(x, sigma=4):
  """Gaussian distance kernel.

  Args:
    x (Tensor): difference between two input features.
    sigma (float): standard deviation of the gaussian kernel.

  Returns:
    The Gaussian kernel for features `x` and standard deviation `sigma`.
  """
  return torch.exp(- (x ** 2).sum(dim=-1) / sigma)