# Path: vmoe/utils.py
# def partialclass(cls: Type[Any], *base_args, **base_kwargs):
#     def __init__(self, *args, **kwargs):
#   def __init__(self, *iterators):
#   def __iter__(self):
#   def __next__(self) -> Tuple[Any, ...]:
# def safe_map(f: Callable[..., Any], *iterables) -> Iterator[Any]:
# def safe_zip(*iterables) -> Iterator[Tuple[Any, ...]]:
# def tree_rngs_split(rngs, num_splits=2):
# def tree_shape_dtype_struct(tree):
#   def fn(x):
#   class _NewClass(cls):
# class SafeZipIteratorError(RuntimeError):
# class SafeZipIterator:
# 
# 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 functools
import time
import cachetools
import flax.core
import flax.struct
import jax
import jax.numpy as jnp
import numpy as np
import tensorflow as tf
from typing import Any, Callable, Dict, Iterable, Mapping, Optional, Sequence, Tuple
from clu import metric_writers
from clu import periodic_actions
from jax.experimental import pjit
from vmoe import utils
from vmoe.data import input_pipeline
    dataset: tf.data.Dataset,
    params: PyTree,
    rng_keys: Sequence[str],
    seed: int = 0,
) -> EvalState:
  """Evaluates a given model on the given dataset."""
  eval_state = EvalState(
      num=np.zeros((), dtype=np.float32),
      sum_correct=np.zeros((), dtype=np.float32),
      sum_loss=np.zeros((), dtype=np.float32),
      rngs=make_rngs(tuple(rng_keys), seed))
  eval_iter = input_pipeline.make_dataset_iterator(dataset)
  for batch in eval_iter:
    eval_state = eval_step_pjit(eval_state, params, batch['image'],
                                batch['labels'], batch['_fake'])
  return jax.tree_map(lambda x: x.block_until_ready(), eval_state)


def evaluate_step(
    state: EvalState,
    params: PyTree,
    images: Array,
    labels: Array,
    fake: Array,
    apply_fn: Callable[..., Any],
    loss_fn: Callable[[Array, Array], Array],
    label_pred_fn: Callable[[Array], Array],
) -> EvalState:
  """Performs one evaluation step, updating the given state."""
  valid = jnp.logical_not(fake)

# Path: vmoe/utils.py
# def partialclass(cls: Type[Any], *base_args, **base_kwargs):
#     def __init__(self, *args, **kwargs):
#   def __init__(self, *iterators):
#   def __iter__(self):
#   def __next__(self) -> Tuple[Any, ...]:
# def safe_map(f: Callable[..., Any], *iterables) -> Iterator[Any]:
# def safe_zip(*iterables) -> Iterator[Tuple[Any, ...]]:
# def tree_rngs_split(rngs, num_splits=2):
# def tree_shape_dtype_struct(tree):
#   def fn(x):
#   class _NewClass(cls):
# class SafeZipIteratorError(RuntimeError):
# class SafeZipIterator:
# 
# 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/evaluate/evaluator.py
import functools
import time
import cachetools
import flax.core
import flax.struct
import jax
import jax.numpy as jnp
import numpy as np
import tensorflow as tf
from typing import Any, Callable, Dict, Iterable, Mapping, Optional, Sequence, Tuple
from clu import metric_writers
from clu import periodic_actions
from jax.experimental import pjit
from vmoe import utils
from vmoe.data import input_pipeline

    dataset: tf.data.Dataset,
    params: PyTree,
    rng_keys: Sequence[str],
    seed: int = 0,
) -> EvalState:
  """Evaluates a given model on the given dataset."""
  eval_state = EvalState(
      num=np.zeros((), dtype=np.float32),
      sum_correct=np.zeros((), dtype=np.float32),
      sum_loss=np.zeros((), dtype=np.float32),
      rngs=make_rngs(tuple(rng_keys), seed))
  eval_iter = input_pipeline.make_dataset_iterator(dataset)
  for batch in eval_iter:
    eval_state = eval_step_pjit(eval_state, params, batch['image'],
                                batch['labels'], batch['_fake'])
  return jax.tree_map(lambda x: x.block_until_ready(), eval_state)


def evaluate_step(
    state: EvalState,
    params: PyTree,
    images: Array,
    labels: Array,
    fake: Array,
    apply_fn: Callable[..., Any],
    loss_fn: Callable[[Array, Array], Array],
    label_pred_fn: Callable[[Array], Array],
) -> EvalState:
  """Performs one evaluation step, updating the given state."""
  valid = jnp.logical_not(fake)

  rngs, next_rngs = utils.tree_rngs_split(state.rngs)