# Path: init2winit/hyperparameters.py
# def build_hparams(model_name,
#                   initializer_name,
#                   dataset_name,
#                   hparam_file,
#                   hparam_overrides):
# 
# Path: init2winit/dataset_lib/datasets.py
# _ALL_DATASETS = {
#     'mnist':
#         _Dataset(small_image_datasets.get_mnist,
#                  small_image_datasets.MNIST_HPARAMS,
#                  small_image_datasets.MNIST_METADATA),
#     'mnist_autoencoder':
#         _Dataset(small_image_datasets.get_mnist_autoencoder,
#                  small_image_datasets.MNIST_AUTOENCODER_HPARAMS,
#                  small_image_datasets.MNIST_AUTOENCODER_METADATA),
#     'fashion_mnist':
#         _Dataset(small_image_datasets.get_fashion_mnist,
#                  small_image_datasets.FASHION_MNIST_HPARAMS,
#                  small_image_datasets.FASHION_MNIST_METADATA),
#     'cifar10':
#         _Dataset(small_image_datasets.get_cifar10,
#                  small_image_datasets.CIFAR10_DEFAULT_HPARAMS,
#                  small_image_datasets.CIFAR10_METADATA),
#     'cifar100':
#         _Dataset(small_image_datasets.get_cifar100,
#                  small_image_datasets.CIFAR100_DEFAULT_HPARAMS,
#                  small_image_datasets.CIFAR100_METADATA),
#     'criteo1tb':
#         _Dataset(criteo_terabyte_input_pipeline.get_criteo1tb,
#                  criteo_terabyte_input_pipeline.CRITEO1TB_DEFAULT_HPARAMS,
#                  criteo_terabyte_input_pipeline.CRITEO1TB_METADATA),
#     'fake':
#         _Dataset(fake_dataset.get_fake, fake_dataset.DEFAULT_HPARAMS,
#                  fake_dataset.METADATA),
#     'imagenet':
#         _Dataset(imagenet_dataset.get_imagenet,
#                  imagenet_dataset.DEFAULT_HPARAMS, imagenet_dataset.METADATA),
#     'translate_wmt':
#         _Dataset(translate_wmt.get_translate_wmt, translate_wmt.DEFAULT_HPARAMS,
#                  translate_wmt.METADATA),
#     'mlperf_imagenet':
#         _Dataset(mlperf_imagenet_dataset.get_mlperf_imagenet,
#                  mlperf_imagenet_dataset.DEFAULT_HPARAMS,
#                  mlperf_imagenet_dataset.METADATA),
#     'svhn_no_extra':
#         _Dataset(small_image_datasets.get_svhn_no_extra,
#                  small_image_datasets.SVHN_NO_EXTRA_DEFAULT_HPARAMS,
#                  small_image_datasets.SVHN_NO_EXTRA_METADATA),
#     'nqm_noise':
#         _Dataset(nqm_noise.get_nqm_noise, nqm_noise.NQM_HPARAMS,
#                  nqm_noise.NQM_METADATA),
#     'ogbg_molpcba':
#         _Dataset(ogbg_molpcba.get_ogbg_molpcba, ogbg_molpcba.DEFAULT_HPARAMS,
#                  ogbg_molpcba.METADATA),
#     'uniref50':
#         _Dataset(proteins.get_uniref, proteins.DEFAULT_HPARAMS,
#                  proteins.METADATA),
# }
# def get_dataset(dataset_name):
# def get_dataset_hparams(dataset_name):
# def get_dataset_meta_data(dataset_name):


import os
import sys
import jax
import tensorflow as tf
from absl import app
from absl import flags
from absl import logging
from init2winit import hyperparameters
from init2winit.dataset_lib import datasets

    if FLAGS.batch_size is None or FLAGS.batch_size <= 0:
      raise ValueError("""FLAGS.batch_size value is invalid,
          expected a positive non-zero integer.""")

    if FLAGS.dataset is None:
      raise ValueError("""FLAGS.dataset value is invalid,
          expected a non-empty string describing dataset name.""")

    batch_size = FLAGS.batch_size
    num_batches = FLAGS.num_batches

    dataset_name = 'lm1b'
    model_name = 'transformer'
    initializer_name = 'noop'

    hparam_overrides = {
        'batch_size': batch_size,
    }

    hps = hyperparameters.build_hparams(
        model_name=model_name,
        initializer_name=initializer_name,
        dataset_name=dataset_name,
        hparam_file=None,
        hparam_overrides=hparam_overrides)

    rng = jax.random.PRNGKey(0)
    rng, data_rng = jax.random.split(rng)


# Path: init2winit/hyperparameters.py
# def build_hparams(model_name,
#                   initializer_name,
#                   dataset_name,
#                   hparam_file,
#                   hparam_overrides):
# 
# Path: init2winit/dataset_lib/datasets.py
# _ALL_DATASETS = {
#     'mnist':
#         _Dataset(small_image_datasets.get_mnist,
#                  small_image_datasets.MNIST_HPARAMS,
#                  small_image_datasets.MNIST_METADATA),
#     'mnist_autoencoder':
#         _Dataset(small_image_datasets.get_mnist_autoencoder,
#                  small_image_datasets.MNIST_AUTOENCODER_HPARAMS,
#                  small_image_datasets.MNIST_AUTOENCODER_METADATA),
#     'fashion_mnist':
#         _Dataset(small_image_datasets.get_fashion_mnist,
#                  small_image_datasets.FASHION_MNIST_HPARAMS,
#                  small_image_datasets.FASHION_MNIST_METADATA),
#     'cifar10':
#         _Dataset(small_image_datasets.get_cifar10,
#                  small_image_datasets.CIFAR10_DEFAULT_HPARAMS,
#                  small_image_datasets.CIFAR10_METADATA),
#     'cifar100':
#         _Dataset(small_image_datasets.get_cifar100,
#                  small_image_datasets.CIFAR100_DEFAULT_HPARAMS,
#                  small_image_datasets.CIFAR100_METADATA),
#     'criteo1tb':
#         _Dataset(criteo_terabyte_input_pipeline.get_criteo1tb,
#                  criteo_terabyte_input_pipeline.CRITEO1TB_DEFAULT_HPARAMS,
#                  criteo_terabyte_input_pipeline.CRITEO1TB_METADATA),
#     'fake':
#         _Dataset(fake_dataset.get_fake, fake_dataset.DEFAULT_HPARAMS,
#                  fake_dataset.METADATA),
#     'imagenet':
#         _Dataset(imagenet_dataset.get_imagenet,
#                  imagenet_dataset.DEFAULT_HPARAMS, imagenet_dataset.METADATA),
#     'translate_wmt':
#         _Dataset(translate_wmt.get_translate_wmt, translate_wmt.DEFAULT_HPARAMS,
#                  translate_wmt.METADATA),
#     'mlperf_imagenet':
#         _Dataset(mlperf_imagenet_dataset.get_mlperf_imagenet,
#                  mlperf_imagenet_dataset.DEFAULT_HPARAMS,
#                  mlperf_imagenet_dataset.METADATA),
#     'svhn_no_extra':
#         _Dataset(small_image_datasets.get_svhn_no_extra,
#                  small_image_datasets.SVHN_NO_EXTRA_DEFAULT_HPARAMS,
#                  small_image_datasets.SVHN_NO_EXTRA_METADATA),
#     'nqm_noise':
#         _Dataset(nqm_noise.get_nqm_noise, nqm_noise.NQM_HPARAMS,
#                  nqm_noise.NQM_METADATA),
#     'ogbg_molpcba':
#         _Dataset(ogbg_molpcba.get_ogbg_molpcba, ogbg_molpcba.DEFAULT_HPARAMS,
#                  ogbg_molpcba.METADATA),
#     'uniref50':
#         _Dataset(proteins.get_uniref, proteins.DEFAULT_HPARAMS,
#                  proteins.METADATA),
# }
# def get_dataset(dataset_name):
# def get_dataset_hparams(dataset_name):
# def get_dataset_meta_data(dataset_name):


# Path: init2winit/tools/inspect_dataset.py
import os
import sys
import jax
import tensorflow as tf
from absl import app
from absl import flags
from absl import logging
from init2winit import hyperparameters
from init2winit.dataset_lib import datasets


    if FLAGS.batch_size is None or FLAGS.batch_size <= 0:
      raise ValueError("""FLAGS.batch_size value is invalid,
          expected a positive non-zero integer.""")

    if FLAGS.dataset is None:
      raise ValueError("""FLAGS.dataset value is invalid,
          expected a non-empty string describing dataset name.""")

    batch_size = FLAGS.batch_size
    num_batches = FLAGS.num_batches

    dataset_name = 'lm1b'
    model_name = 'transformer'
    initializer_name = 'noop'

    hparam_overrides = {
        'batch_size': batch_size,
    }

    hps = hyperparameters.build_hparams(
        model_name=model_name,
        initializer_name=initializer_name,
        dataset_name=dataset_name,
        hparam_file=None,
        hparam_overrides=hparam_overrides)

    rng = jax.random.PRNGKey(0)
    rng, data_rng = jax.random.split(rng)


    dataset = datasets.get_dataset(FLAGS.dataset)(data_rng, batch_size,