# 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
      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)
    filenames = filenames.shard(self._dataset_num_shards, self._dataset_index)

    def make_dataset(ds_index, repeat=True):
      ds = filenames.shard(self._parallelism, ds_index)
      if repeat:
        ds = ds.repeat()
      ds = ds.interleave(
          tf.data.TextLineDataset,
          cycle_length=16,
          block_length=batch_size // 8,
          num_parallel_calls=8,
          deterministic=False)
      ds = ds.batch(batch_size, drop_remainder=True)
      ds = ds.map(_parse_example_fn, num_parallel_calls=16)
      return ds


# 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

      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)
    filenames = filenames.shard(self._dataset_num_shards, self._dataset_index)

    def make_dataset(ds_index, repeat=True):
      ds = filenames.shard(self._parallelism, ds_index)
      if repeat:
        ds = ds.repeat()
      ds = ds.interleave(
          tf.data.TextLineDataset,
          cycle_length=16,
          block_length=batch_size // 8,
          num_parallel_calls=8,
          deterministic=False)
      ds = ds.batch(batch_size, drop_remainder=True)
      ds = ds.map(_parse_example_fn, num_parallel_calls=16)
      return ds


    ds_indices = tf.data.Dataset.range(self._parallelism)