def clip_eta(eta, ord, eps):
  """
  PyTorch implementation of the clip_eta in utils_tf.

  :param eta: Tensor
  :param ord: np.inf, 1, or 2
  :param eps: float
  """
  if ord not in [np.inf, 1, 2]:
    raise ValueError('ord must be np.inf, 1, or 2.')

  avoid_zero_div = torch.tensor(1e-12, dtype=eta.dtype, device=eta.device)
  reduc_ind = list(range(1, len(eta.size())))
  if ord == np.inf:
    eta = torch.clamp(eta, -eps, eps)
  else:
    if ord == 1:
      # TODO
      # raise NotImplementedError("L1 clip is not implemented.")
      norm = torch.max(
          avoid_zero_div,
          torch.sum(torch.abs(eta), dim=reduc_ind, keepdim=True)
      )
    elif ord == 2:
      norm = torch.sqrt(torch.max(
          avoid_zero_div,
          torch.sum(eta ** 2, dim=reduc_ind, keepdim=True)
      ))
    factor = torch.min(
        torch.tensor(1., dtype=eta.dtype, device=eta.device),
        eps / norm
        )
    eta *= factor
  return eta
PyTorch implementation of the clip_eta in utils_tf.

  :param eta: Tensor
  :param ord: np.inf, 1, or 2
  :param eps: float