# 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/lm1b_input_pipeline_v2.py
# AUTOTUNE = tf.data.experimental.AUTOTUNE
# class NormalizeFeatureNamesOp:
#   def __init__(self, ds_info: tfds.core.DatasetInfo):
#   def __call__(self, features: Features) -> Features:
# def get_raw_dataset(dataset_builder: tfds.core.DatasetBuilder,
#                     split: str) -> tf.data.Dataset:
# def pack_dataset(dataset: tf.data.Dataset,
#                  key2length: Union[int, Dict[str, int]],
#                  keys: Optional[List[str]] = None) -> tf.data.Dataset:
#   def my_fn(x):
# def _pack_with_tf_ops(dataset: tf.data.Dataset, keys: List[str],
#                       key2length: Dict[str, int]) -> tf.data.Dataset:
#   def write_packed_example(partial, outputs):
#   def map_fn(x):
#     def body_fn(i, partial, outputs):
#       def false_fn():
#       def true_fn():
# def preprocess_data(dataset,
#                     shuffle=True,
#                     num_epochs=1,
#                     pack_examples=False,
#                     shuffle_buffer_size=1024,
#                     max_length=512,
#                     batch_size=256,
#                     drop_remainder=True,
#                     prefetch_size=AUTOTUNE,
#                     shuffle_seed=None):
#   def length_filter(max_len):
#     def filter_fn(x):
# def get_lm1b_datasets(hps, per_host_batch_size, per_host_eval_batch_size,
#                       shuffle_rng):
# 
# 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 itertools
import jax
import jax.numpy as jnp
from init2winit.dataset_lib import data_utils
from init2winit.dataset_lib import lm1b_input_pipeline_v2
from init2winit.dataset_lib.data_utils import Dataset
from ml_collections.config_dict import config_dict
  if per_host_eval_batch_size % n_devices != 0:
    raise ValueError(
        'n_devices={} must divide per_host_eval_batch_size={}.'.format(
            n_devices, per_host_eval_batch_size))

  train_ds, eval_ds = lm1b_input_pipeline_v2.get_lm1b_datasets(
      hps, per_host_batch_size, per_host_eval_batch_size, shuffle_rng)

  def train_iterator_fn():
    for batch in iter(train_ds):
      yield _batch_to_dict(batch)

  def eval_train_epoch(num_batches=None):
    eval_train_iter = iter(train_ds)
    for batch in itertools.islice(eval_train_iter, num_batches):
      yield _batch_to_dict(batch)

  def valid_epoch(num_batches=None):
    valid_iter = iter(eval_ds)
    for batch in itertools.islice(valid_iter, num_batches):
      yield _batch_to_dict(batch)

  # pylint: disable=unreachable
  def test_epoch(*args, **kwargs):
    del args
    del kwargs
    return
    yield  # This yield is needed to make this a valid (null) iterator.

  # pylint: enable=unreachable

# 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/lm1b_input_pipeline_v2.py
# AUTOTUNE = tf.data.experimental.AUTOTUNE
# class NormalizeFeatureNamesOp:
#   def __init__(self, ds_info: tfds.core.DatasetInfo):
#   def __call__(self, features: Features) -> Features:
# def get_raw_dataset(dataset_builder: tfds.core.DatasetBuilder,
#                     split: str) -> tf.data.Dataset:
# def pack_dataset(dataset: tf.data.Dataset,
#                  key2length: Union[int, Dict[str, int]],
#                  keys: Optional[List[str]] = None) -> tf.data.Dataset:
#   def my_fn(x):
# def _pack_with_tf_ops(dataset: tf.data.Dataset, keys: List[str],
#                       key2length: Dict[str, int]) -> tf.data.Dataset:
#   def write_packed_example(partial, outputs):
#   def map_fn(x):
#     def body_fn(i, partial, outputs):
#       def false_fn():
#       def true_fn():
# def preprocess_data(dataset,
#                     shuffle=True,
#                     num_epochs=1,
#                     pack_examples=False,
#                     shuffle_buffer_size=1024,
#                     max_length=512,
#                     batch_size=256,
#                     drop_remainder=True,
#                     prefetch_size=AUTOTUNE,
#                     shuffle_seed=None):
#   def length_filter(max_len):
#     def filter_fn(x):
# def get_lm1b_datasets(hps, per_host_batch_size, per_host_eval_batch_size,
#                       shuffle_rng):
# 
# 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/lm1b_v2.py
import itertools
import jax
import jax.numpy as jnp
from init2winit.dataset_lib import data_utils
from init2winit.dataset_lib import lm1b_input_pipeline_v2
from init2winit.dataset_lib.data_utils import Dataset
from ml_collections.config_dict import config_dict

  if per_host_eval_batch_size % n_devices != 0:
    raise ValueError(
        'n_devices={} must divide per_host_eval_batch_size={}.'.format(
            n_devices, per_host_eval_batch_size))

  train_ds, eval_ds = lm1b_input_pipeline_v2.get_lm1b_datasets(
      hps, per_host_batch_size, per_host_eval_batch_size, shuffle_rng)

  def train_iterator_fn():
    for batch in iter(train_ds):
      yield _batch_to_dict(batch)

  def eval_train_epoch(num_batches=None):
    eval_train_iter = iter(train_ds)
    for batch in itertools.islice(eval_train_iter, num_batches):
      yield _batch_to_dict(batch)

  def valid_epoch(num_batches=None):
    valid_iter = iter(eval_ds)
    for batch in itertools.islice(valid_iter, num_batches):
      yield _batch_to_dict(batch)

  # pylint: disable=unreachable
  def test_epoch(*args, **kwargs):
    del args
    del kwargs
    return
    yield  # This yield is needed to make this a valid (null) iterator.

  # pylint: enable=unreachable

  return Dataset(train_iterator_fn, eval_train_epoch, valid_epoch, test_epoch)