# Path: init2winit/dataset_lib/datasets.py
# def get_dataset(dataset_name):
#   """Maps dataset name to a dataset_builder."""
#   try:
#     return _ALL_DATASETS[dataset_name].getter
#   except KeyError:
#     raise ValueError('Unrecognized dataset: {}'.format(dataset_name))
# 
# Path: init2winit/dataset_lib/datasets.py
# def get_dataset_hparams(dataset_name):
#   """Maps dataset name to default_hps."""
#   try:
#     hparams = _ALL_DATASETS[dataset_name].hparams
#     # TODO(mbadura): Refactor to explicitly support different input specs
#     if 'input_shape' not in hparams or hparams.input_shape is None:
#       if 'input_edge_shape' in hparams and 'input_node_shape' in hparams:
#         hparams.input_shape = (hparams.input_node_shape,
#                                hparams.input_edge_shape)
#       elif dataset_name == 'lm1b':
#         max_len = max(hparams.max_target_length, hparams.max_eval_target_length)
#         hparams.input_shape = (max_len,)
#       elif dataset_name == 'translate_wmt':
#         max_len = max(hparams.max_target_length, hparams.max_eval_target_length,
#                       hparams.max_predict_length)
#         hparams.input_shape = [(max_len,), (max_len,)]
#       else:
#         raise ValueError(
#             'Undefined input shape for dataset: {}'.format(dataset_name))
#     return hparams
#   except KeyError:
#     raise ValueError('Unrecognized dataset: {}'.format(dataset_name))


import itertools
import jax.random
import numpy as np
import tensorflow.compat.v1 as tf
import tensorflow_datasets as tfds
from init2winit.dataset_lib.datasets import get_dataset
from init2winit.dataset_lib.datasets import get_dataset_hparams
          labels
  }


def _as_dataset(*args, **kwargs):
  """Creates a mock TFDS graphs dataset."""
  del args, kwargs

  def get_iter():
    return (
        _make_graph(num_nodes, num_edges, labels)
        for num_nodes, num_edges, labels in zip(NUMS_NODES, NUMS_EDGES, LABELS))

  return tf.data.Dataset.from_generator(
      get_iter,
      output_signature={
          'edge_feat': tf.TensorSpec(shape=(None, 3), dtype=np.float32),
          'edge_index': tf.TensorSpec(shape=(None, 2), dtype=np.int64),
          'labels': tf.TensorSpec(shape=(NUM_LABELS,), dtype=np.float32),
          'node_feat': tf.TensorSpec(shape=(None, 9), dtype=np.float32),
          'num_edges': tf.TensorSpec(shape=(1,), dtype=np.int64),
          'num_nodes': tf.TensorSpec(shape=(1,), dtype=np.int64),
      })


def _get_dataset(shuffle_seed):
  """Loads the ogbg-molpcba dataset using mock data."""
  with tfds.testing.mock_data(as_dataset_fn=_as_dataset):
    ds = 'ogbg_molpcba'
    dataset_builder = get_dataset(ds)

# Path: init2winit/dataset_lib/datasets.py
# def get_dataset(dataset_name):
#   """Maps dataset name to a dataset_builder."""
#   try:
#     return _ALL_DATASETS[dataset_name].getter
#   except KeyError:
#     raise ValueError('Unrecognized dataset: {}'.format(dataset_name))
# 
# Path: init2winit/dataset_lib/datasets.py
# def get_dataset_hparams(dataset_name):
#   """Maps dataset name to default_hps."""
#   try:
#     hparams = _ALL_DATASETS[dataset_name].hparams
#     # TODO(mbadura): Refactor to explicitly support different input specs
#     if 'input_shape' not in hparams or hparams.input_shape is None:
#       if 'input_edge_shape' in hparams and 'input_node_shape' in hparams:
#         hparams.input_shape = (hparams.input_node_shape,
#                                hparams.input_edge_shape)
#       elif dataset_name == 'lm1b':
#         max_len = max(hparams.max_target_length, hparams.max_eval_target_length)
#         hparams.input_shape = (max_len,)
#       elif dataset_name == 'translate_wmt':
#         max_len = max(hparams.max_target_length, hparams.max_eval_target_length,
#                       hparams.max_predict_length)
#         hparams.input_shape = [(max_len,), (max_len,)]
#       else:
#         raise ValueError(
#             'Undefined input shape for dataset: {}'.format(dataset_name))
#     return hparams
#   except KeyError:
#     raise ValueError('Unrecognized dataset: {}'.format(dataset_name))


# Path: init2winit/dataset_lib/test_ogbg_molpcba.py
import itertools
import jax.random
import numpy as np
import tensorflow.compat.v1 as tf
import tensorflow_datasets as tfds
from init2winit.dataset_lib.datasets import get_dataset
from init2winit.dataset_lib.datasets import get_dataset_hparams

          labels
  }


def _as_dataset(*args, **kwargs):
  """Creates a mock TFDS graphs dataset."""
  del args, kwargs

  def get_iter():
    return (
        _make_graph(num_nodes, num_edges, labels)
        for num_nodes, num_edges, labels in zip(NUMS_NODES, NUMS_EDGES, LABELS))

  return tf.data.Dataset.from_generator(
      get_iter,
      output_signature={
          'edge_feat': tf.TensorSpec(shape=(None, 3), dtype=np.float32),
          'edge_index': tf.TensorSpec(shape=(None, 2), dtype=np.int64),
          'labels': tf.TensorSpec(shape=(NUM_LABELS,), dtype=np.float32),
          'node_feat': tf.TensorSpec(shape=(None, 9), dtype=np.float32),
          'num_edges': tf.TensorSpec(shape=(1,), dtype=np.int64),
          'num_nodes': tf.TensorSpec(shape=(1,), dtype=np.int64),
      })


def _get_dataset(shuffle_seed):
  """Loads the ogbg-molpcba dataset using mock data."""
  with tfds.testing.mock_data(as_dataset_fn=_as_dataset):
    ds = 'ogbg_molpcba'
    dataset_builder = get_dataset(ds)

    hps = get_dataset_hparams(ds)