# 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

  The Maximum Mean Discrepancy (MMD) is a measure of the distance between the
  distributions of prediction scores on two groups of examples. The metric
  guarantees that the result is 0 if and only if the two distributions it is
  comparing are exactly the same.

  The `membership` input indicates with a numerical value whether
  each example is part of the sensitive group with a numerical value. This
  currently only supports hard membership of `0.0` or `1.0`.

  For more details, see the
  [paper](http://papers.nips.cc/paper/3110-a-kernel-method-for-the-two-sample-problem.pdf).
  """
  # pyformat: enable

  def __init__(self,
               kernel="gaussian",
               predictions_transform=None,
               name: Optional[str] = None,
               enable_summary_histogram: Optional[bool] = True):
    """Initialize an instance of MMDLoss."""
    super(MMDLoss, self).__init__(
        predictions_transform=predictions_transform,
        predictions_kernel=kernel,
        name=name or "mmd_loss",
        enable_summary_histogram=enable_summary_histogram)

  def _preprocess(self,
                  membership: types.TensorType,
                  predictions: 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/mmd_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


  The Maximum Mean Discrepancy (MMD) is a measure of the distance between the
  distributions of prediction scores on two groups of examples. The metric
  guarantees that the result is 0 if and only if the two distributions it is
  comparing are exactly the same.

  The `membership` input indicates with a numerical value whether
  each example is part of the sensitive group with a numerical value. This
  currently only supports hard membership of `0.0` or `1.0`.

  For more details, see the
  [paper](http://papers.nips.cc/paper/3110-a-kernel-method-for-the-two-sample-problem.pdf).
  """
  # pyformat: enable

  def __init__(self,
               kernel="gaussian",
               predictions_transform=None,
               name: Optional[str] = None,
               enable_summary_histogram: Optional[bool] = True):
    """Initialize an instance of MMDLoss."""
    super(MMDLoss, self).__init__(
        predictions_transform=predictions_transform,
        predictions_kernel=kernel,
        name=name or "mmd_loss",
        enable_summary_histogram=enable_summary_histogram)

  def _preprocess(self,
                  membership: types.TensorType,
                  predictions: types.TensorType,

                  sample_weight: Optional[types.TensorType] = None):