# 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/mt_pipeline.py
# AUTOTUNE = tf.data.AUTOTUNE
# def get_user_defined_symbols(ds_info: tfds.core.DatasetInfo,
#                              reverse_translation: bool):
#   def __init__(self, ds_info: tfds.core.DatasetInfo, reverse_translation: bool):
#   def __call__(self, features: Features) -> Features:
#   def __init__(self, ds_info: tfds.core.DatasetInfo,
#                reverse_translation: bool):
#   def __call__(self, features):
#   def __init__(self, weight):
#   def __call__(self, features):
# def maybe_pad_batch(batch,
#                     desired_batch_size,
#                     mask_key='targets'):
#   def zero_pad(ar, pad_axis):
# def get_raw_dataset(dataset_builder: tfds.core.DatasetBuilder,
#                     split: str,
#                     *,
#                     reverse_translation: bool = False,
#                     add_language_token: bool = False) -> 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 get_sampled_dataset(ds_builders,
#                         split: str,
#                         rates: List[int],
#                         reverse_translation: bool,
#                         add_language_token: bool,
#                         loss_weights: List[float],
#                         is_training: bool,
#                         sample_seed: int,
#                         shuffle_seed: int,
#                         shuffle_buffer_size: int = 1024):
#   def _shuffle_repeat(dataset, shuffle_seed: int,
#                       shuffle_buffer_size):
# def preprocess_wmt_data(dataset,
#                         pack_examples: bool = True,
#                         max_length: int = 100,
#                         batch_size: int = 256,
#                         prefetch_size: int = AUTOTUNE):
#   def length_filter(max_len):
#     def filter_fn(x):
# def get_wmt_datasets(config: config_dict.ConfigDict,
#                      *,
#                      shuffle_seed: int,
#                      sample_seed: int,
#                      n_devices: int,
#                      per_host_batch_size: int,
#                      per_host_eval_batch_size: int,
#                      vocab_path: Optional[str] = None):
# class NormalizeFeatureNamesOp:
# class TaskTokenOp:
# class ImportanceSamplingOp:
# 
# 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
from init2winit.dataset_lib import data_utils
from init2winit.dataset_lib import mt_pipeline
from init2winit.dataset_lib.data_utils import Dataset
from ml_collections.config_dict import config_dict

  vocab_path = hps.vocab_path
  train_ds, eval_ds, predict_ds = mt_pipeline.get_wmt_datasets(
      hps,
      shuffle_seed=shuffle_rng[0],
      sample_seed=shuffle_rng[1],
      n_devices=jax.local_device_count(),
      per_host_batch_size=per_host_batch_size,
      per_host_eval_batch_size=per_host_eval_batch_size,
      vocab_path=vocab_path)

  def train_iterator_fn():
    for batch in iter(train_ds):
      yield mt_pipeline.maybe_pad_batch(
          data_utils.tf_to_numpy(batch),
          per_host_batch_size,
          mask_key='targets')

  def eval_train_epoch(num_batches=None):
    eval_train_iter = iter(train_ds)
    for batch in itertools.islice(eval_train_iter, num_batches):
      yield mt_pipeline.maybe_pad_batch(
          data_utils.tf_to_numpy(batch),
          per_host_batch_size,
          mask_key='targets')

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

# 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/mt_pipeline.py
# AUTOTUNE = tf.data.AUTOTUNE
# def get_user_defined_symbols(ds_info: tfds.core.DatasetInfo,
#                              reverse_translation: bool):
#   def __init__(self, ds_info: tfds.core.DatasetInfo, reverse_translation: bool):
#   def __call__(self, features: Features) -> Features:
#   def __init__(self, ds_info: tfds.core.DatasetInfo,
#                reverse_translation: bool):
#   def __call__(self, features):
#   def __init__(self, weight):
#   def __call__(self, features):
# def maybe_pad_batch(batch,
#                     desired_batch_size,
#                     mask_key='targets'):
#   def zero_pad(ar, pad_axis):
# def get_raw_dataset(dataset_builder: tfds.core.DatasetBuilder,
#                     split: str,
#                     *,
#                     reverse_translation: bool = False,
#                     add_language_token: bool = False) -> 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 get_sampled_dataset(ds_builders,
#                         split: str,
#                         rates: List[int],
#                         reverse_translation: bool,
#                         add_language_token: bool,
#                         loss_weights: List[float],
#                         is_training: bool,
#                         sample_seed: int,
#                         shuffle_seed: int,
#                         shuffle_buffer_size: int = 1024):
#   def _shuffle_repeat(dataset, shuffle_seed: int,
#                       shuffle_buffer_size):
# def preprocess_wmt_data(dataset,
#                         pack_examples: bool = True,
#                         max_length: int = 100,
#                         batch_size: int = 256,
#                         prefetch_size: int = AUTOTUNE):
#   def length_filter(max_len):
#     def filter_fn(x):
# def get_wmt_datasets(config: config_dict.ConfigDict,
#                      *,
#                      shuffle_seed: int,
#                      sample_seed: int,
#                      n_devices: int,
#                      per_host_batch_size: int,
#                      per_host_eval_batch_size: int,
#                      vocab_path: Optional[str] = None):
# class NormalizeFeatureNamesOp:
# class TaskTokenOp:
# class ImportanceSamplingOp:
# 
# 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/translate_wmt.py
import itertools
import jax
from init2winit.dataset_lib import data_utils
from init2winit.dataset_lib import mt_pipeline
from init2winit.dataset_lib.data_utils import Dataset
from ml_collections.config_dict import config_dict


  vocab_path = hps.vocab_path
  train_ds, eval_ds, predict_ds = mt_pipeline.get_wmt_datasets(
      hps,
      shuffle_seed=shuffle_rng[0],
      sample_seed=shuffle_rng[1],
      n_devices=jax.local_device_count(),
      per_host_batch_size=per_host_batch_size,
      per_host_eval_batch_size=per_host_eval_batch_size,
      vocab_path=vocab_path)

  def train_iterator_fn():
    for batch in iter(train_ds):
      yield mt_pipeline.maybe_pad_batch(
          data_utils.tf_to_numpy(batch),
          per_host_batch_size,
          mask_key='targets')

  def eval_train_epoch(num_batches=None):
    eval_train_iter = iter(train_ds)
    for batch in itertools.islice(eval_train_iter, num_batches):
      yield mt_pipeline.maybe_pad_batch(
          data_utils.tf_to_numpy(batch),
          per_host_batch_size,
          mask_key='targets')

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

          data_utils.tf_to_numpy(batch),