# 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




def _get_batch(tensors, batch_size, batch_num):
  if isinstance(tensors, dict):
    return {k: _get_batch(v, batch_size, batch_num) for k, v in tensors.items()}

  if isinstance(tensors, tf.SparseTensor):
    total_examples = tensors.dense_shape[0]
  else:
    total_examples = len(tensors)
  start_ind = (batch_size * batch_num) % total_examples
  end_ind = start_ind + batch_size
  # Double tensor to enable repeating inputs.
  if isinstance(tensors, tf.SparseTensor):
    return tf.sparse.slice(tensors, [start_ind, 0], [batch_size, 2])
  else:
    return tensors[start_ind:end_ind]


def _get_min_diff_batch(sensitive_tensors, nonsensitive_tensors,
                        sensitive_batch_size, nonsensitive_batch_size,
                        batch_num):
  sensitive_batch = _get_batch(sensitive_tensors, sensitive_batch_size,
                               batch_num)
  nonsensitive_batch = _get_batch(nonsensitive_tensors, nonsensitive_batch_size,
                                  batch_num)
  if isinstance(sensitive_batch, dict):
    return {

# 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





def _get_batch(tensors, batch_size, batch_num):
  if isinstance(tensors, dict):
    return {k: _get_batch(v, batch_size, batch_num) for k, v in tensors.items()}

  if isinstance(tensors, tf.SparseTensor):
    total_examples = tensors.dense_shape[0]
  else:
    total_examples = len(tensors)
  start_ind = (batch_size * batch_num) % total_examples
  end_ind = start_ind + batch_size
  # Double tensor to enable repeating inputs.
  if isinstance(tensors, tf.SparseTensor):
    return tf.sparse.slice(tensors, [start_ind, 0], [batch_size, 2])
  else:
    return tensors[start_ind:end_ind]


def _get_min_diff_batch(sensitive_tensors, nonsensitive_tensors,
                        sensitive_batch_size, nonsensitive_batch_size,
                        batch_num):
  sensitive_batch = _get_batch(sensitive_tensors, sensitive_batch_size,
                               batch_num)
  nonsensitive_batch = _get_batch(nonsensitive_tensors, nonsensitive_batch_size,
                                  batch_num)
  if isinstance(sensitive_batch, dict):
    return {

        key: input_utils._tensor_concat(sensitive_batch[key],