# Path: tensorflow_model_remediation/min_diff/losses/absolute_correlation_loss.py
# _EPSILON = 1.0e-7
# class AbsoluteCorrelationLoss(base_loss.MinDiffLoss):
#   def __init__(self, name: Optional[str] = None,
#                enable_summary_histogram: Optional[bool] = True):
#   def call(
#       self,
#       sensitive_group_labels: types.TensorType,
#       y_pred: types.TensorType,
#       sample_weight: Optional[types.TensorType] = None) -> types.TensorType:
# 
# Path: tensorflow_model_remediation/min_diff/losses/base_loss.py
# class MinDiffLoss(tf.keras.losses.Loss, abc.ABC):
#   def __init__(self,
#                membership_transform=None,
#                predictions_transform=None,
#                membership_kernel=None,
#                predictions_kernel=None,
#                name: Optional[str] = None,
#                enable_summary_histogram: Optional[bool] = True):
#   def __call__(self,
#                membership: types.TensorType,
#                predictions: types.TensorType,
#                sample_weight: Optional[types.TensorType] = None):
#   def _preprocess_inputs(
#       self,
#       membership: types.TensorType,
#       predictions: types.TensorType,
#       sample_weight: Optional[types.TensorType] = None
#   ) -> Tuple[types.TensorType, types.TensorType, types.TensorType]:
#   def _apply_kernels(
#       self, membership: types.TensorType, predictions: types.TensorType
#   ) -> Tuple[types.TensorType, types.TensorType]:
#   def call(self,
#            membership: types.TensorType,
#            predictions: types.TensorType,
#            sample_weight: Optional[types.TensorType] = None):
#   def _serialize_config(self, config):
#     def _serialize_value(key, value):
#   def get_config(self):
#   def _deserialize_config(cls, config):
#     def _deserialize_value(key, value):
#   def from_config(cls, config):
# def _validate_transform(transform: types.TensorTransformType,
#                         var_name: str) -> None:
# def _to_snake_case(name):
# 
# Path: tensorflow_model_remediation/min_diff/losses/mmd_loss.py
# class MMDLoss(base_loss.MinDiffLoss):
#   def __init__(self,
#                kernel="gaussian",
#                predictions_transform=None,
#                name: Optional[str] = None,
#                enable_summary_histogram: Optional[bool] = True):
#   def _preprocess(self,
#                   membership: types.TensorType,
#                   predictions: types.TensorType,
#                   sample_weight: Optional[types.TensorType] = None):
#   def _calculate_mean(self, predictions_kernel: types.TensorType,
#                       normed_weights: types.TensorType,
#                       pos_mask: types.TensorType, neg_mask):
#   def call(self,
#            membership: types.TensorType,
#            predictions: types.TensorType,
#            sample_weight: Optional[types.TensorType] = None):
#   def from_config(cls, config):


from typing import Union
from tensorflow_model_remediation.min_diff.losses import absolute_correlation_loss as abs_corr_loss
from tensorflow_model_remediation.min_diff.losses import base_loss
from tensorflow_model_remediation.min_diff.losses import mmd_loss


_register_loss_names(abs_corr_loss.AbsoluteCorrelationLoss,
                     ['abs_corr', 'absolute_correlation'])
_register_loss_names(mmd_loss.MMDLoss, ['mmd', 'maximum_mean_discrepancy'])


def _get_loss(loss: Union[base_loss.MinDiffLoss, str],
              loss_var_name: str = 'loss') -> base_loss.MinDiffLoss:
  """Returns a `losses.MinDiffLoss` instance corresponding to `loss`.

  If `loss` is an instance of `losses.MinDiffLoss` then it is returned
  directly. If `loss` is a string it must be an accepted loss name. A
  value of `None` is also accepted and simply returns `None`.

  Args:
    loss: loss instance. Can be `None`, a string or an instance of
      `losses.MinDiffLoss`.
    loss_var_name: Name of the loss variable. This is only used for error
      messaging.

  Returns:
    A `MinDiffLoss` instance.

  Raises:
    ValueError: If `loss` is an unrecognized string.
    TypeError: If `loss` is not an instance of `losses.MinDiffLoss` or a string.
  """
  if loss is None:
    return None

# Path: tensorflow_model_remediation/min_diff/losses/absolute_correlation_loss.py
# _EPSILON = 1.0e-7
# class AbsoluteCorrelationLoss(base_loss.MinDiffLoss):
#   def __init__(self, name: Optional[str] = None,
#                enable_summary_histogram: Optional[bool] = True):
#   def call(
#       self,
#       sensitive_group_labels: types.TensorType,
#       y_pred: types.TensorType,
#       sample_weight: Optional[types.TensorType] = None) -> types.TensorType:
# 
# Path: tensorflow_model_remediation/min_diff/losses/base_loss.py
# class MinDiffLoss(tf.keras.losses.Loss, abc.ABC):
#   def __init__(self,
#                membership_transform=None,
#                predictions_transform=None,
#                membership_kernel=None,
#                predictions_kernel=None,
#                name: Optional[str] = None,
#                enable_summary_histogram: Optional[bool] = True):
#   def __call__(self,
#                membership: types.TensorType,
#                predictions: types.TensorType,
#                sample_weight: Optional[types.TensorType] = None):
#   def _preprocess_inputs(
#       self,
#       membership: types.TensorType,
#       predictions: types.TensorType,
#       sample_weight: Optional[types.TensorType] = None
#   ) -> Tuple[types.TensorType, types.TensorType, types.TensorType]:
#   def _apply_kernels(
#       self, membership: types.TensorType, predictions: types.TensorType
#   ) -> Tuple[types.TensorType, types.TensorType]:
#   def call(self,
#            membership: types.TensorType,
#            predictions: types.TensorType,
#            sample_weight: Optional[types.TensorType] = None):
#   def _serialize_config(self, config):
#     def _serialize_value(key, value):
#   def get_config(self):
#   def _deserialize_config(cls, config):
#     def _deserialize_value(key, value):
#   def from_config(cls, config):
# def _validate_transform(transform: types.TensorTransformType,
#                         var_name: str) -> None:
# def _to_snake_case(name):
# 
# Path: tensorflow_model_remediation/min_diff/losses/mmd_loss.py
# class MMDLoss(base_loss.MinDiffLoss):
#   def __init__(self,
#                kernel="gaussian",
#                predictions_transform=None,
#                name: Optional[str] = None,
#                enable_summary_histogram: Optional[bool] = True):
#   def _preprocess(self,
#                   membership: types.TensorType,
#                   predictions: types.TensorType,
#                   sample_weight: Optional[types.TensorType] = None):
#   def _calculate_mean(self, predictions_kernel: types.TensorType,
#                       normed_weights: types.TensorType,
#                       pos_mask: types.TensorType, neg_mask):
#   def call(self,
#            membership: types.TensorType,
#            predictions: types.TensorType,
#            sample_weight: Optional[types.TensorType] = None):
#   def from_config(cls, config):


# Path: tensorflow_model_remediation/min_diff/losses/loss_utils.py
from typing import Union
from tensorflow_model_remediation.min_diff.losses import absolute_correlation_loss as abs_corr_loss
from tensorflow_model_remediation.min_diff.losses import base_loss
from tensorflow_model_remediation.min_diff.losses import mmd_loss



_register_loss_names(abs_corr_loss.AbsoluteCorrelationLoss,
                     ['abs_corr', 'absolute_correlation'])
_register_loss_names(mmd_loss.MMDLoss, ['mmd', 'maximum_mean_discrepancy'])


def _get_loss(loss: Union[base_loss.MinDiffLoss, str],
              loss_var_name: str = 'loss') -> base_loss.MinDiffLoss:
  """Returns a `losses.MinDiffLoss` instance corresponding to `loss`.

  If `loss` is an instance of `losses.MinDiffLoss` then it is returned
  directly. If `loss` is a string it must be an accepted loss name. A
  value of `None` is also accepted and simply returns `None`.

  Args:
    loss: loss instance. Can be `None`, a string or an instance of
      `losses.MinDiffLoss`.
    loss_var_name: Name of the loss variable. This is only used for error
      messaging.

  Returns:
    A `MinDiffLoss` instance.

  Raises:
    ValueError: If `loss` is an unrecognized string.
    TypeError: If `loss` is not an instance of `losses.MinDiffLoss` or a string.
  """
  if loss is None:
    return None

  if isinstance(loss, base_loss.MinDiffLoss):