# Path: tensorflow_model_remediation/min_diff/keras/utils/input_utils.py
# class MinDiffPackedInputs(
#     collections.namedtuple("MinDiffPackedInputs",
#                            ["original_inputs", "min_diff_data"])):
# def pack_min_diff_data(original_dataset: tf.data.Dataset,
#                        sensitive_group_dataset=None,
#                        nonsensitive_group_dataset=None,
#                        min_diff_dataset=None) -> tf.data.Dataset:
#   def _map_fn(original_batch, min_diff_batch):
# def _pack_as_original(original_batch, x, y, w):
# def _tensor_concat(t1, t2):
# def build_min_diff_dataset(sensitive_group_dataset,
#                            nonsensitive_group_dataset) -> tf.data.Dataset:
#   def _build_single_batch(single_sensitive_batch, single_nonsensitive_batch):
#   def _map_fn(sensitive_batch, nonsensitive_batch):
# def unpack_original_inputs(inputs):
# def unpack_min_diff_data(inputs):


import copy
import tensorflow as tf
from tensorflow_model_remediation.min_diff.keras.utils import input_utils
    with self.assertRaisesRegex(
        ValueError, "x component structure.*sensitive_group_dataset.*does not "
        "match.*nonsensitive_group_dataset(.|\n)*don't have the same nested "
        "structure"):
      _ = input_utils.pack_min_diff_data(original_dataset, sensitive_dataset,
                                         nonsensitive_dataset)


class UnpackDataTest(tf.test.TestCase):

  def testUnpackOriginal(self):
    # Tensor.
    tensor = tf.fill([3, 4], 1.3)
    packed_inputs = input_utils.MinDiffPackedInputs(tensor, None)
    unpacked_tensor = input_utils.unpack_original_inputs(packed_inputs)
    self.assertIs(unpacked_tensor, tensor)

    # Dict of Tensors.
    tensors = {"f1": tf.fill([1, 2], 2), "f2": tf.fill([4, 5], "a")}
    packed_inputs = input_utils.MinDiffPackedInputs(tensors, None)
    unpacked_tensors = input_utils.unpack_original_inputs(packed_inputs)
    self.assertIs(unpacked_tensors, tensors)

    # Arbitrary object.
    obj = set(["a", "b", "c"])
    packed_inputs = input_utils.MinDiffPackedInputs(obj, None)
    unpacked_obj = input_utils.unpack_original_inputs(packed_inputs)
    self.assertIs(unpacked_obj, obj)

    # None.

# Path: tensorflow_model_remediation/min_diff/keras/utils/input_utils.py
# class MinDiffPackedInputs(
#     collections.namedtuple("MinDiffPackedInputs",
#                            ["original_inputs", "min_diff_data"])):
# def pack_min_diff_data(original_dataset: tf.data.Dataset,
#                        sensitive_group_dataset=None,
#                        nonsensitive_group_dataset=None,
#                        min_diff_dataset=None) -> tf.data.Dataset:
#   def _map_fn(original_batch, min_diff_batch):
# def _pack_as_original(original_batch, x, y, w):
# def _tensor_concat(t1, t2):
# def build_min_diff_dataset(sensitive_group_dataset,
#                            nonsensitive_group_dataset) -> tf.data.Dataset:
#   def _build_single_batch(single_sensitive_batch, single_nonsensitive_batch):
#   def _map_fn(sensitive_batch, nonsensitive_batch):
# def unpack_original_inputs(inputs):
# def unpack_min_diff_data(inputs):


# Path: tensorflow_model_remediation/min_diff/keras/utils/input_utils_test.py
import copy
import tensorflow as tf
from tensorflow_model_remediation.min_diff.keras.utils import input_utils

    with self.assertRaisesRegex(
        ValueError, "x component structure.*sensitive_group_dataset.*does not "
        "match.*nonsensitive_group_dataset(.|\n)*don't have the same nested "
        "structure"):
      _ = input_utils.pack_min_diff_data(original_dataset, sensitive_dataset,
                                         nonsensitive_dataset)


class UnpackDataTest(tf.test.TestCase):

  def testUnpackOriginal(self):
    # Tensor.
    tensor = tf.fill([3, 4], 1.3)
    packed_inputs = input_utils.MinDiffPackedInputs(tensor, None)
    unpacked_tensor = input_utils.unpack_original_inputs(packed_inputs)
    self.assertIs(unpacked_tensor, tensor)

    # Dict of Tensors.
    tensors = {"f1": tf.fill([1, 2], 2), "f2": tf.fill([4, 5], "a")}
    packed_inputs = input_utils.MinDiffPackedInputs(tensors, None)
    unpacked_tensors = input_utils.unpack_original_inputs(packed_inputs)
    self.assertIs(unpacked_tensors, tensors)

    # Arbitrary object.
    obj = set(["a", "b", "c"])
    packed_inputs = input_utils.MinDiffPackedInputs(obj, None)
    unpacked_obj = input_utils.unpack_original_inputs(packed_inputs)
    self.assertIs(unpacked_obj, obj)

    # None.

    packed_inputs = input_utils.MinDiffPackedInputs(None, None)