# Path: vmoe/data/input_pipeline.py
# DEFAULT_SHUFFLE_BUFFER = 50_000
# def _get_dataset_builder(name, data_dir, manual_dir, try_gcs):
# def get_datasets(
#     config: ml_collections.ConfigDict) -> Dict[str, tf.data.Dataset]:
# def get_data_from_tfds(*,
#                        variant: str,
#                        name: str,
#                        split: str,
#                        batch_size: int,
#                        process: str,
#                        cache: Optional[str] = None,
#                        data_dir: Optional[str] = None,
#                        manual_dir: Optional[str] = None,
#                        prefetch: Optional[Union[int, str]] = None,
#                        shuffle_buffer: int = DEFAULT_SHUFFLE_BUFFER,
#                        shuffle_seed: Optional[int] = None,
#                        try_gcs: bool = False) -> tf.data.Dataset:
# def get_data_num_examples(config: ml_collections.ConfigDict) -> int:
# def get_data_process_fn(process_str: str) -> Callable[[Data], Data]:
#   def _process_fn(data: Data):
# def get_data_range(builder, split: str, process_id: int, process_count: int):
# def make_dataset_iterator(dataset: tf.data.Dataset) -> Iterator[Dict[str, Any]]:
#   def to_numpy(data):
# def _parse_process_op_str(string_to_parse):
#   def _get_func_name(expr):
#   def _get_func_args_and_kwargs(call):


import collections
import ml_collections
import numpy as np
import tensorflow as tf
from unittest import mock
from absl.testing import absltest
from vmoe.data import input_pipeline
  def test_get_data_process_fn_wrong_dataset_error(self):
    """Tests that an exception is raised if the dataset is not made of dicts."""
    process_fn = input_pipeline.get_data_process_fn('copy("foo", "bar")')
    dataset = tf.data.Dataset.from_tensors(tf.convert_to_tensor([0.0]))
    with self.assertRaises(TypeError):
      dataset = dataset.map(process_fn)
      _ = next(iter(dataset))

  def test_get_data_range(self):
    _SplitInfo = collections.namedtuple('SplitInfo', ('num_examples',))
    builder = mock.MagicMock()
    builder.info.splits = {'x': _SplitInfo(num_examples=11)}
    data_ranges = [
        input_pipeline.get_data_range(builder, 'x', i, 4) for i in range(4)
    ]
    self.assertTupleEqual(data_ranges[0], ('x', 0, 3, False))
    self.assertTupleEqual(data_ranges[1], ('x', 3, 6, False))
    self.assertTupleEqual(data_ranges[2], ('x', 6, 9, False))
    self.assertTupleEqual(data_ranges[3], ('x', 9, 11, True))

  def test_get_data_num_examples(self):
    builder = mock.MagicMock()
    _SplitInfo = collections.namedtuple('SplitInfo', ('num_examples',))
    builder.info.splits = {'bar': _SplitInfo(num_examples=3)}
    with mock.patch.object(
        input_pipeline.tfds, 'builder', return_value=builder):
      config = ml_collections.ConfigDict()
      config.name = 'foo'
      config.split = 'bar'
      self.assertEqual(

# Path: vmoe/data/input_pipeline.py
# DEFAULT_SHUFFLE_BUFFER = 50_000
# def _get_dataset_builder(name, data_dir, manual_dir, try_gcs):
# def get_datasets(
#     config: ml_collections.ConfigDict) -> Dict[str, tf.data.Dataset]:
# def get_data_from_tfds(*,
#                        variant: str,
#                        name: str,
#                        split: str,
#                        batch_size: int,
#                        process: str,
#                        cache: Optional[str] = None,
#                        data_dir: Optional[str] = None,
#                        manual_dir: Optional[str] = None,
#                        prefetch: Optional[Union[int, str]] = None,
#                        shuffle_buffer: int = DEFAULT_SHUFFLE_BUFFER,
#                        shuffle_seed: Optional[int] = None,
#                        try_gcs: bool = False) -> tf.data.Dataset:
# def get_data_num_examples(config: ml_collections.ConfigDict) -> int:
# def get_data_process_fn(process_str: str) -> Callable[[Data], Data]:
#   def _process_fn(data: Data):
# def get_data_range(builder, split: str, process_id: int, process_count: int):
# def make_dataset_iterator(dataset: tf.data.Dataset) -> Iterator[Dict[str, Any]]:
#   def to_numpy(data):
# def _parse_process_op_str(string_to_parse):
#   def _get_func_name(expr):
#   def _get_func_args_and_kwargs(call):


# Path: vmoe/data/input_pipeline_test.py
import collections
import ml_collections
import numpy as np
import tensorflow as tf
from unittest import mock
from absl.testing import absltest
from vmoe.data import input_pipeline

  def test_get_data_process_fn_wrong_dataset_error(self):
    """Tests that an exception is raised if the dataset is not made of dicts."""
    process_fn = input_pipeline.get_data_process_fn('copy("foo", "bar")')
    dataset = tf.data.Dataset.from_tensors(tf.convert_to_tensor([0.0]))
    with self.assertRaises(TypeError):
      dataset = dataset.map(process_fn)
      _ = next(iter(dataset))

  def test_get_data_range(self):
    _SplitInfo = collections.namedtuple('SplitInfo', ('num_examples',))
    builder = mock.MagicMock()
    builder.info.splits = {'x': _SplitInfo(num_examples=11)}
    data_ranges = [
        input_pipeline.get_data_range(builder, 'x', i, 4) for i in range(4)
    ]
    self.assertTupleEqual(data_ranges[0], ('x', 0, 3, False))
    self.assertTupleEqual(data_ranges[1], ('x', 3, 6, False))
    self.assertTupleEqual(data_ranges[2], ('x', 6, 9, False))
    self.assertTupleEqual(data_ranges[3], ('x', 9, 11, True))

  def test_get_data_num_examples(self):
    builder = mock.MagicMock()
    _SplitInfo = collections.namedtuple('SplitInfo', ('num_examples',))
    builder.info.splits = {'bar': _SplitInfo(num_examples=3)}
    with mock.patch.object(
        input_pipeline.tfds, 'builder', return_value=builder):
      config = ml_collections.ConfigDict()
      config.name = 'foo'
      config.split = 'bar'
      self.assertEqual(

          input_pipeline.get_data_num_examples(config=config), 3)