# 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/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 jraph
import numpy as np
import tensorflow_datasets as tfds
from init2winit.dataset_lib import data_utils
from init2winit.dataset_lib.data_utils import Dataset
from ml_collections.config_dict import config_dict
                  shuffle_seed=None,
                  shuffle_buffer_size=None):
  """Loads a dataset split from TFDS."""
  if should_shuffle:
    assert shuffle_seed is not None and shuffle_buffer_size is not None
    file_shuffle_seed, dataset_shuffle_seed = jax.random.split(shuffle_seed)
    file_shuffle_seed = file_shuffle_seed[0]
    dataset_shuffle_seed = dataset_shuffle_seed[0]
  else:
    file_shuffle_seed = None
    dataset_shuffle_seed = None

  read_config = tfds.ReadConfig(
      add_tfds_id=True, shuffle_seed=file_shuffle_seed)
  dataset = tfds.load(
      'ogbg_molpcba',
      split=split,
      shuffle_files=should_shuffle,
      read_config=read_config)

  if should_shuffle:
    dataset = dataset.shuffle(
        seed=dataset_shuffle_seed, buffer_size=shuffle_buffer_size)
    dataset = dataset.repeat()

  return dataset


def _to_jraph(example):
  """Converts an example graph to jraph.GraphsTuple."""

# 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/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/ogbg_molpcba.py
import itertools
import jax
import jraph
import numpy as np
import tensorflow_datasets as tfds
from init2winit.dataset_lib import data_utils
from init2winit.dataset_lib.data_utils import Dataset
from ml_collections.config_dict import config_dict

                  shuffle_seed=None,
                  shuffle_buffer_size=None):
  """Loads a dataset split from TFDS."""
  if should_shuffle:
    assert shuffle_seed is not None and shuffle_buffer_size is not None
    file_shuffle_seed, dataset_shuffle_seed = jax.random.split(shuffle_seed)
    file_shuffle_seed = file_shuffle_seed[0]
    dataset_shuffle_seed = dataset_shuffle_seed[0]
  else:
    file_shuffle_seed = None
    dataset_shuffle_seed = None

  read_config = tfds.ReadConfig(
      add_tfds_id=True, shuffle_seed=file_shuffle_seed)
  dataset = tfds.load(
      'ogbg_molpcba',
      split=split,
      shuffle_files=should_shuffle,
      read_config=read_config)

  if should_shuffle:
    dataset = dataset.shuffle(
        seed=dataset_shuffle_seed, buffer_size=shuffle_buffer_size)
    dataset = dataset.repeat()

  return dataset


def _to_jraph(example):
  """Converts an example graph to jraph.GraphsTuple."""

  example = data_utils.tf_to_numpy(example)