import torch


def one_hot(y, K, smooth_eps = 0):  # pylint: disable=invalid-name
  """One-hot encodes a tensor with optional label smoothing.

  Args:
    y: A tensor containing the ground-truth labels of shape (N,), i.e. one label
      for each element in the batch.
    K: The number of classes.
    smooth_eps: Label smoothing factor in [0, 1] range.

  Returns:
    A one-hot encoded tensor.
  """
  assert 0 <= smooth_eps <= 1
  assert y.ndim == 1, "Label tensor must be rank 1."
  y_hot = torch.eye(K)[y] * (1 - smooth_eps) + (smooth_eps / (K - 1))
  return y_hot.to(y.device)