# Path: tensorflow_model_remediation/common/types.py

# 
# 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):


from typing import Optional
from tensorflow_model_remediation.common import types
from tensorflow_model_remediation.min_diff.losses import base_loss
import tensorflow as tf
  """Absolute correlation between predictions on two groups of examples.

  Arguments:
    name: Name used for logging or tracking. Defaults to
      `'absolute_correlation_loss'`.
    enable_summary_histogram: Optional bool indicating if `tf.summary.histogram`
      should be included within the loss. Defaults to True.

  Absolute correlation measures how correlated predictions are with membership
  (regardless of direction). The metric guarantees that the result is 0 if and
  only if the two distributions it is comparing are indistinguishable.

  The `sensitive_group_labels` input is used to determine whether each example
  is part of the sensitive group. This currently only supports hard membership
  of `0.0` or `1.0`.

  For more details, see the [paper](https://arxiv.org/abs/1901.04562).
  """
  # pyformat: enable

  def __init__(self, name: Optional[str] = None,
               enable_summary_histogram: Optional[bool] = True):
    """Initialize Loss."""
    super(AbsoluteCorrelationLoss,
          self).__init__(name=name or 'absolute_correlation_loss',
                         enable_summary_histogram=enable_summary_histogram)

  def call(
      self,
      sensitive_group_labels: types.TensorType,

# Path: tensorflow_model_remediation/common/types.py

# 
# 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/absolute_correlation_loss.py
from typing import Optional
from tensorflow_model_remediation.common import types
from tensorflow_model_remediation.min_diff.losses import base_loss
import tensorflow as tf

  """Absolute correlation between predictions on two groups of examples.

  Arguments:
    name: Name used for logging or tracking. Defaults to
      `'absolute_correlation_loss'`.
    enable_summary_histogram: Optional bool indicating if `tf.summary.histogram`
      should be included within the loss. Defaults to True.

  Absolute correlation measures how correlated predictions are with membership
  (regardless of direction). The metric guarantees that the result is 0 if and
  only if the two distributions it is comparing are indistinguishable.

  The `sensitive_group_labels` input is used to determine whether each example
  is part of the sensitive group. This currently only supports hard membership
  of `0.0` or `1.0`.

  For more details, see the [paper](https://arxiv.org/abs/1901.04562).
  """
  # pyformat: enable

  def __init__(self, name: Optional[str] = None,
               enable_summary_histogram: Optional[bool] = True):
    """Initialize Loss."""
    super(AbsoluteCorrelationLoss,
          self).__init__(name=name or 'absolute_correlation_loss',
                         enable_summary_histogram=enable_summary_histogram)

  def call(
      self,
      sensitive_group_labels: types.TensorType,

      y_pred: types.TensorType,