content: # Copyright (c) Meta Platforms, Inc. and affiliates.
  #
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  
  import abc
  
  import torch
  
  _LOSS_EPS = 1e-20
  
  
  class RobustLoss(abc.ABC):
      @classmethod
      def evaluate(cls, x: torch.Tensor, log_radius: torch.Tensor) -> torch.Tensor:
          return cls._evaluate_impl(x, log_radius.exp())
  
      @classmethod
      def linearize(cls, x: torch.Tensor, log_radius: torch.Tensor) -> torch.Tensor:
          return cls._linearize_impl(x, log_radius.exp())
  
      @staticmethod
      @abc.abstractmethod
      def _evaluate_impl(x: torch.Tensor, radius: torch.Tensor) -> torch.Tensor:
          pass
  
      @staticmethod
      @abc.abstractmethod
      def _linearize_impl(x: torch.Tensor, raidus: torch.Tensor) -> torch.Tensor:
          pass
  
  
  class WelschLoss(RobustLoss):
      @staticmethod
      def _evaluate_impl(x: torch.Tensor, radius: torch.Tensor) -> torch.Tensor:
          return radius - radius * torch.exp(-x / (radius + _LOSS_EPS))
  
      @staticmethod
      def _linearize_impl(x: torch.Tensor, radius: torch.Tensor) -> torch.Tensor:
          return torch.exp(-x / (radius + _LOSS_EPS))
  
  
  class HuberLoss(RobustLoss):
      @staticmethod
      def _evaluate_impl(x: torch.Tensor, radius: torch.Tensor) -> torch.Tensor:
          return torch.where(
              x > radius, 2 * torch.sqrt(radius * x.max(radius) + _LOSS_EPS) - radius, x
          )
  
      @staticmethod
      def _linearize_impl(x: torch.Tensor, radius: torch.Tensor) -> torch.Tensor:
          return torch.sqrt(radius / torch.max(x, radius) + _LOSS_EPS)
  
  
  class HingeLoss(RobustLoss):
      @staticmethod
      def _evaluate_impl(x: torch.Tensor, radius: torch.Tensor) -> torch.Tensor:
          return torch.where(x > radius, torch.sqrt(x) - torch.sqrt(radius), _LOSS_EPS)
  
      @staticmethod
      def _linearize_impl(x: torch.Tensor, radius: torch.Tensor) -> torch.Tensor:
          return torch.where(x > radius, 1.0 / (2 * torch.sqrt(x) + _LOSS_EPS), 0.0)
  
  
  class GNCRobustLoss(RobustLoss, abc.ABC):
      @classmethod
      def evaluate(  # type: ignore
          cls, x: torch.Tensor, log_radius: torch.Tensor, mu: torch.Tensor
      ) -> torch.Tensor:
          return cls._evaluate_impl(x, log_radius.exp(), mu)
  
      @classmethod
      def linearize(  # type: ignore
          cls, x: torch.Tensor, log_radius: torch.Tensor, mu: torch.Tensor
      ) -> torch.Tensor:
          return cls._linearize_impl(x, log_radius.exp(), mu)
  
      @staticmethod
      @abc.abstractmethod
      def _evaluate_impl(  # type: ignore
          x: torch.Tensor, radius: torch.Tensor, mu: torch.Tensor
      ) -> torch.Tensor:
          pass
  
      @staticmethod
      @abc.abstractmethod
      def _linearize_impl(  # type: ignore
          x: torch.Tensor, raidus: torch.Tensor, mu: torch.Tensor
      ) -> torch.Tensor:
          pass
  
  
  # Geman-McClure loss is one common robust loss that can be adaptively tuned
  # with a continous control parameter `mu`, hence very suited for GNC.
  # More context on adaptive robust loss at https://arxiv.org/pdf/1701.03077.pdf
  class GemanMcClureLoss(GNCRobustLoss):
      @staticmethod
      def _evaluate_impl(  # type: ignore
          x: torch.Tensor, radius: torch.Tensor, mu: torch.Tensor
      ) -> torch.Tensor:
          """
          x: squared loss (i.e., r^2)
          radius: scale of quadratic bowl (also squared)
          mu: GNC control parameter ranging from 1 (full Geman McClure) to
              positive infinity (full Quadratic)
          """
          return mu * radius * x / (mu * radius + x + _LOSS_EPS)
  
      @staticmethod
      def _linearize_impl(  # type: ignore
          x: torch.Tensor, radius: torch.Tensor, mu: torch.Tensor
      ) -> torch.Tensor:
          return (mu * radius) ** 2 / ((mu * radius + x) ** 2 + _LOSS_EPS)
  