# Path: init2winit/dataset_lib/data_utils.py
# def iterator_as_numpy(iterator):
# def image_iterator(data,
#                    rescale,
#                    output_shape,
#                    is_one_hot,
#                    autoencoder,
#                    shuffle_rng=None,
#                    augment_fn=None,
#                    include_example_keys=False):
# def maybe_pad_batch(batch,
#                     desired_batch_size,
#                     data_format=None,
#                     mask_key=None):
#   def zero_pad(ar, pad_axis):
# def shard(batch, n_devices=None):
#   def _shard_array(array):
# def tf_to_numpy(tfds_data):


import functools
import itertools
import os
import tensorflow as tf
import tensorflow_datasets as tfds
from absl import logging
from init2winit.dataset_lib.data_utils import Dataset
from ml_collections.config_dict import config_dict
    def _parse_example_fn(example):
      """Parser function for pre-processed Criteo TSV records."""
      label_defaults = [[0.0]]
      int_defaults = [[0.0] for _ in range(self._num_dense_features)]
      categorical_defaults = [[0] for _ in range(len(self._vocab_sizes))]
      record_defaults = label_defaults + int_defaults + categorical_defaults
      fields = tf.io.decode_csv(
          example, record_defaults, field_delim='\t', na_value='-1')

      num_labels = 1
      num_dense = len(int_defaults)
      features = {}
      features['targets'] = tf.reshape(fields[0], [batch_size, 1])

      int_features = []
      for idx in range(num_dense):
        int_features.append(fields[idx + num_labels])
      int_features = tf.stack(int_features, axis=1)

      cat_features = []
      for idx in range(len(self._vocab_sizes)):
        cat_features.append(
            tf.cast(fields[idx + num_dense + num_labels], dtype=tf.int32))
      cat_features = tf.cast(
          tf.stack(cat_features, axis=1), dtype=int_features.dtype)
      features['inputs'] = tf.concat([int_features, cat_features], axis=1)
      features['weights'] = tf.ones(
          shape=(features['inputs'].shape[0],), dtype=features['inputs'].dtype)
      return features


# Path: init2winit/dataset_lib/data_utils.py
# def iterator_as_numpy(iterator):
# def image_iterator(data,
#                    rescale,
#                    output_shape,
#                    is_one_hot,
#                    autoencoder,
#                    shuffle_rng=None,
#                    augment_fn=None,
#                    include_example_keys=False):
# def maybe_pad_batch(batch,
#                     desired_batch_size,
#                     data_format=None,
#                     mask_key=None):
#   def zero_pad(ar, pad_axis):
# def shard(batch, n_devices=None):
#   def _shard_array(array):
# def tf_to_numpy(tfds_data):


# Path: init2winit/dataset_lib/criteo_terabyte_input_pipeline.py
import functools
import itertools
import os
import tensorflow as tf
import tensorflow_datasets as tfds
from absl import logging
from init2winit.dataset_lib.data_utils import Dataset
from ml_collections.config_dict import config_dict

    def _parse_example_fn(example):
      """Parser function for pre-processed Criteo TSV records."""
      label_defaults = [[0.0]]
      int_defaults = [[0.0] for _ in range(self._num_dense_features)]
      categorical_defaults = [[0] for _ in range(len(self._vocab_sizes))]
      record_defaults = label_defaults + int_defaults + categorical_defaults
      fields = tf.io.decode_csv(
          example, record_defaults, field_delim='\t', na_value='-1')

      num_labels = 1
      num_dense = len(int_defaults)
      features = {}
      features['targets'] = tf.reshape(fields[0], [batch_size, 1])

      int_features = []
      for idx in range(num_dense):
        int_features.append(fields[idx + num_labels])
      int_features = tf.stack(int_features, axis=1)

      cat_features = []
      for idx in range(len(self._vocab_sizes)):
        cat_features.append(
            tf.cast(fields[idx + num_dense + num_labels], dtype=tf.int32))
      cat_features = tf.cast(
          tf.stack(cat_features, axis=1), dtype=int_features.dtype)
      features['inputs'] = tf.concat([int_features, cat_features], axis=1)
      features['weights'] = tf.ones(
          shape=(features['inputs'].shape[0],), dtype=features['inputs'].dtype)
      return features


    filenames = tf.data.Dataset.list_files(self._file_path, shuffle=False)