# Path: vmoe/evaluate/evaluator.py
# class EvalState(flax.struct.PyTreeNode):
# class EvaluateMultipleDatasets(periodic_actions.PeriodicCallback):
#   def update(self, num, correct, loss, rngs):
#   def __init__(
#       self,
#       *,
#       apply_fn: Callable[..., Any],
#       loss_fn: Callable[[Array, Array], Array],
#       label_pred_fn: Callable[[Array], Array],
#       params_axis_resources: PyTree,
#       input_axis_resources: PartitionSpec,
#       datasets: Mapping[str, tf.data.Dataset],
#       metric_writer: metric_writers.MetricWriter,
#       rng_keys: Sequence[str],
#       seed: int = 0,
#       every_steps: Optional[int] = None,
#       every_secs: Optional[float] = None,
#       on_steps: Optional[Iterable[int]] = None,
#       report_progress: Optional[periodic_actions.ReportProgress] = None,
#       report_progress_name: str = 'eval'):
#   def _make_callback_fn(cls, *, apply_fn, loss_fn, label_pred_fn,
#                         params_axis_resources, input_axis_resources, datasets,
#                         metric_writer, rng_keys, seed, report_progress,
#                         report_progress_name):
#     def compile_for_dataset(name, params, train_step):
#     def callback_fn(step: int, t: Optional[float], params: PyTree):
# def evaluate_dataset(
#     *,
#     eval_step_pjit: EvalStepPjitFn,
#     dataset: tf.data.Dataset,
#     params: PyTree,
#     rng_keys: Sequence[str],
#     seed: int = 0,
# ) -> EvalState:
# 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:
# def make_eval_step_pjit(
#     params_axis_resources: PyTree,
#     input_axis_resources: PartitionSpec,
#     apply_fn: Callable[..., Any],
#     loss_fn: Callable[[Array, Array], Array],
#     label_pred_fn: Callable[[Array], Array],
#     rng_keys: Sequence[str],
# ) -> EvalStepPjitFn:
# def make_rngs(rng_keys: Tuple[str, ...], seed: int = 0) -> Dict[str, PRNGKey]:
#   def _make_rngs():


from unittest import mock
from absl.testing import absltest
from jax.experimental import maps
from vmoe.evaluate import evaluator
import jax
import numpy as np
import tensorflow as tf

  def test_evaluate_dataset(self):
    # Create random test dataset.
    dataset, expected_eval_state = self._create_dataset_and_expected_state()

    eval_step_pjit = evaluator.make_eval_step_pjit(
        params_axis_resources={},  # Model has no parameters.
        input_axis_resources=PartitionSpec('d',),
        apply_fn=self._apply_fn,
        loss_fn=self._loss_fn,
        label_pred_fn=self._label_pred_fn,
        rng_keys=[])
    # Run the evaluation.
    with maps.mesh(np.asarray(jax.local_devices()), ('d',)):
      eval_state = evaluator.evaluate_dataset(
          eval_step_pjit=eval_step_pjit,
          dataset=dataset,
          params={},  # Model has no parameters.
          rng_keys=[])
    # Check that the values of the EvalState are correct.
    self.assertAlmostEqual(eval_state.num, expected_eval_state.num)
    self.assertAlmostEqual(eval_state.sum_correct,
                           expected_eval_state.sum_correct)
    self.assertAlmostEqual(eval_state.sum_loss, expected_eval_state.sum_loss,
                           places=4)

  def test_evaluate_multiple_datasets(self):
    dataset1, _ = self._create_dataset_and_expected_state()
    dataset2, _ = self._create_dataset_and_expected_state()
    datasets = {'dataset1': dataset1, 'dataset2': dataset2}

# Path: vmoe/evaluate/evaluator.py
# class EvalState(flax.struct.PyTreeNode):
# class EvaluateMultipleDatasets(periodic_actions.PeriodicCallback):
#   def update(self, num, correct, loss, rngs):
#   def __init__(
#       self,
#       *,
#       apply_fn: Callable[..., Any],
#       loss_fn: Callable[[Array, Array], Array],
#       label_pred_fn: Callable[[Array], Array],
#       params_axis_resources: PyTree,
#       input_axis_resources: PartitionSpec,
#       datasets: Mapping[str, tf.data.Dataset],
#       metric_writer: metric_writers.MetricWriter,
#       rng_keys: Sequence[str],
#       seed: int = 0,
#       every_steps: Optional[int] = None,
#       every_secs: Optional[float] = None,
#       on_steps: Optional[Iterable[int]] = None,
#       report_progress: Optional[periodic_actions.ReportProgress] = None,
#       report_progress_name: str = 'eval'):
#   def _make_callback_fn(cls, *, apply_fn, loss_fn, label_pred_fn,
#                         params_axis_resources, input_axis_resources, datasets,
#                         metric_writer, rng_keys, seed, report_progress,
#                         report_progress_name):
#     def compile_for_dataset(name, params, train_step):
#     def callback_fn(step: int, t: Optional[float], params: PyTree):
# def evaluate_dataset(
#     *,
#     eval_step_pjit: EvalStepPjitFn,
#     dataset: tf.data.Dataset,
#     params: PyTree,
#     rng_keys: Sequence[str],
#     seed: int = 0,
# ) -> EvalState:
# 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:
# def make_eval_step_pjit(
#     params_axis_resources: PyTree,
#     input_axis_resources: PartitionSpec,
#     apply_fn: Callable[..., Any],
#     loss_fn: Callable[[Array, Array], Array],
#     label_pred_fn: Callable[[Array], Array],
#     rng_keys: Sequence[str],
# ) -> EvalStepPjitFn:
# def make_rngs(rng_keys: Tuple[str, ...], seed: int = 0) -> Dict[str, PRNGKey]:
#   def _make_rngs():


# Path: vmoe/evaluate/evaluator_test.py
from unittest import mock
from absl.testing import absltest
from jax.experimental import maps
from vmoe.evaluate import evaluator
import jax
import numpy as np
import tensorflow as tf


  def test_evaluate_dataset(self):
    # Create random test dataset.
    dataset, expected_eval_state = self._create_dataset_and_expected_state()

    eval_step_pjit = evaluator.make_eval_step_pjit(
        params_axis_resources={},  # Model has no parameters.
        input_axis_resources=PartitionSpec('d',),
        apply_fn=self._apply_fn,
        loss_fn=self._loss_fn,
        label_pred_fn=self._label_pred_fn,
        rng_keys=[])
    # Run the evaluation.
    with maps.mesh(np.asarray(jax.local_devices()), ('d',)):
      eval_state = evaluator.evaluate_dataset(
          eval_step_pjit=eval_step_pjit,
          dataset=dataset,
          params={},  # Model has no parameters.
          rng_keys=[])
    # Check that the values of the EvalState are correct.
    self.assertAlmostEqual(eval_state.num, expected_eval_state.num)
    self.assertAlmostEqual(eval_state.sum_correct,
                           expected_eval_state.sum_correct)
    self.assertAlmostEqual(eval_state.sum_loss, expected_eval_state.sum_loss,
                           places=4)

  def test_evaluate_multiple_datasets(self):
    dataset1, _ = self._create_dataset_and_expected_state()
    dataset2, _ = self._create_dataset_and_expected_state()
    datasets = {'dataset1': dataset1, 'dataset2': dataset2}

    metric_writer = mock.MagicMock(evaluator.metric_writers.MetricWriter)