# Path: vmoe/train/trainer.py
# def create_checkpoint_hook(*, workdir: str, progress_hook: ReportProgress,
#                            train_state_axis_resources: TrainStateAxisResources,
#                            train_steps: int,
#                            **kwargs) -> PeriodicCheckpointSaver:
# def create_evaluation_hook(*, writer: metric_writers.MetricWriter,
#                            progress_hook: ReportProgress,
#                            datasets: Mapping[str, Dataset],
#                            apply_fn: Callable[..., Any],
#                            loss_fn: Callable[[Array, Array], Array],
#                            params_axis_resources: PyTree,
#                            input_axis_resources: PartitionSpec, first_step: int,
#                            train_steps: int, extra_rng_keys: Sequence[str],
#                            **kwargs) -> evaluator.EvaluateMultipleDatasets:
# def create_flax_model(*, config: Dict[str, Any],
#                       deterministic: bool) -> nn.Module:
# def create_profile_hook(*, workdir: str, **kwargs):
# def create_progress_hook(*, writer: metric_writers.MetricWriter,
#                          first_step: int, train_steps: int,
#                          **kwargs) -> ReportProgress:
# def create_train_state(
#     *,
#     initialize_fn: TrainStateInitFromScratchFn,
#     axis_resources: TrainStateAxisResources,
#     rngs: Mapping[str, PRNGKey],
#     mesh: Optional[Mesh] = None,
# ) -> TrainState:
# def create_train_state_initialize_from_scratch_fn(
#     *,
#     model: nn.Module,
#     optimizer_config: Dict[str, Any],
#     input_shape: Tuple[int, ...],
#     input_axis_resources: PartitionSpec,
#     train_steps: int) -> TrainStateInitFromScratchFn:
#   def initialize(rngs: Mapping[str, PRNGKey]):
# def restore_or_create_train_state(
#     *,
#     prefix: str,
#     initialize_fn: TrainStateInitFromScratchFn,
#     axis_resources: TrainState,
#     rngs: Mapping[str, PRNGKey],
#     mesh: Optional[Mesh] = None,
#     thread_pool: Optional[ThreadPool] = None,
# ) -> TrainState:
# def get_loss_fn(name: str, **kwargs):
#     def train_loss_fn(logits, labels):
# def get_train_steps_and_epochs(
#     *,
#     train_batch_size: int,
#     train_examples: int,
#     train_steps: Optional[int] = None,
#     train_epochs: Optional[Union[int, float]] = None,
# ) -> Tuple[int, float]:
# def initialize_train_state_from_checkpoint(
#     *,
#     train_state: TrainState,
#     axis_resources: TrainStateAxisResources,
#     name: str,
#     **kwargs) -> TrainState:
# def parse_partition_spec(spec) -> PartitionSpec:
# def train_step(
#     state: TrainState,
#     images: Array,
#     labels: Array,
#     loss_fn: Callable[[Array, Array], Array],
# ) -> Tuple[TrainState, Mapping[str, Any]]:
#   def compute_grads_and_metrics(params):
# def train_and_evaluate(config: ml_collections.ConfigDict, workdir: str):
# def tree_axis_resources_from_regexes(
#     *,
#     tree: PyTree,
#     axis_resources_regexes: AxisResourcesRegexes,
# ) -> PyTree:
#   def search_partition_spec(key: str) -> PartitionSpec:


from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
from vmoe.train import trainer
import chex
import flax.linen as nn
import jax
import jax.numpy as jnp
import ml_collections
import numpy as np
import tensorflow as tf
                                 eval_expected_shape):
    name = 'ensemble_softmax_xent'
    kwargs = {'ensemble_size': ensemble_size}
    train_loss_fn, eval_loss_fn, label_pred_fn = trainer.get_loss_fn(
        name, **kwargs)
    key = jax.random.PRNGKey(0)
    logits = jax.random.normal(key, (8, 16))
    train_labels = jax.random.uniform(key, (8, 16))
    eval_labels = jax.random.uniform(key, (8 // ensemble_size, 16))

    train_loss = train_loss_fn(logits, train_labels)
    eval_loss = eval_loss_fn(logits, eval_labels)

    self.assertEqual(train_loss.shape, train_expected_shape)
    self.assertEqual(eval_loss.shape, eval_expected_shape)
    self.assertEqual(label_pred_fn(logits).shape, eval_expected_shape)
    # When ensemble_size is 1, we are back in the standard case where the train
    # and eval losses are identical.
    if ensemble_size == 1:
      self.assertSequenceEqual(list(train_loss), list(eval_loss))


class GetTrainStepsAndEpochsTest(parameterized.TestCase):

  @parameterized.named_parameters(
      ('_train_epochs', {'train_epochs': 1}, (10, 1.0)),
      ('_train_steps', {'train_steps': 11}, (11, 1.1)),
  )
  def test(self, kwargs, expected):
    self.assertEqual(

# Path: vmoe/train/trainer.py
# def create_checkpoint_hook(*, workdir: str, progress_hook: ReportProgress,
#                            train_state_axis_resources: TrainStateAxisResources,
#                            train_steps: int,
#                            **kwargs) -> PeriodicCheckpointSaver:
# def create_evaluation_hook(*, writer: metric_writers.MetricWriter,
#                            progress_hook: ReportProgress,
#                            datasets: Mapping[str, Dataset],
#                            apply_fn: Callable[..., Any],
#                            loss_fn: Callable[[Array, Array], Array],
#                            params_axis_resources: PyTree,
#                            input_axis_resources: PartitionSpec, first_step: int,
#                            train_steps: int, extra_rng_keys: Sequence[str],
#                            **kwargs) -> evaluator.EvaluateMultipleDatasets:
# def create_flax_model(*, config: Dict[str, Any],
#                       deterministic: bool) -> nn.Module:
# def create_profile_hook(*, workdir: str, **kwargs):
# def create_progress_hook(*, writer: metric_writers.MetricWriter,
#                          first_step: int, train_steps: int,
#                          **kwargs) -> ReportProgress:
# def create_train_state(
#     *,
#     initialize_fn: TrainStateInitFromScratchFn,
#     axis_resources: TrainStateAxisResources,
#     rngs: Mapping[str, PRNGKey],
#     mesh: Optional[Mesh] = None,
# ) -> TrainState:
# def create_train_state_initialize_from_scratch_fn(
#     *,
#     model: nn.Module,
#     optimizer_config: Dict[str, Any],
#     input_shape: Tuple[int, ...],
#     input_axis_resources: PartitionSpec,
#     train_steps: int) -> TrainStateInitFromScratchFn:
#   def initialize(rngs: Mapping[str, PRNGKey]):
# def restore_or_create_train_state(
#     *,
#     prefix: str,
#     initialize_fn: TrainStateInitFromScratchFn,
#     axis_resources: TrainState,
#     rngs: Mapping[str, PRNGKey],
#     mesh: Optional[Mesh] = None,
#     thread_pool: Optional[ThreadPool] = None,
# ) -> TrainState:
# def get_loss_fn(name: str, **kwargs):
#     def train_loss_fn(logits, labels):
# def get_train_steps_and_epochs(
#     *,
#     train_batch_size: int,
#     train_examples: int,
#     train_steps: Optional[int] = None,
#     train_epochs: Optional[Union[int, float]] = None,
# ) -> Tuple[int, float]:
# def initialize_train_state_from_checkpoint(
#     *,
#     train_state: TrainState,
#     axis_resources: TrainStateAxisResources,
#     name: str,
#     **kwargs) -> TrainState:
# def parse_partition_spec(spec) -> PartitionSpec:
# def train_step(
#     state: TrainState,
#     images: Array,
#     labels: Array,
#     loss_fn: Callable[[Array, Array], Array],
# ) -> Tuple[TrainState, Mapping[str, Any]]:
#   def compute_grads_and_metrics(params):
# def train_and_evaluate(config: ml_collections.ConfigDict, workdir: str):
# def tree_axis_resources_from_regexes(
#     *,
#     tree: PyTree,
#     axis_resources_regexes: AxisResourcesRegexes,
# ) -> PyTree:
#   def search_partition_spec(key: str) -> PartitionSpec:


# Path: vmoe/train/trainer_test.py
from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
from vmoe.train import trainer
import chex
import flax.linen as nn
import jax
import jax.numpy as jnp
import ml_collections
import numpy as np
import tensorflow as tf

                                 eval_expected_shape):
    name = 'ensemble_softmax_xent'
    kwargs = {'ensemble_size': ensemble_size}
    train_loss_fn, eval_loss_fn, label_pred_fn = trainer.get_loss_fn(
        name, **kwargs)
    key = jax.random.PRNGKey(0)
    logits = jax.random.normal(key, (8, 16))
    train_labels = jax.random.uniform(key, (8, 16))
    eval_labels = jax.random.uniform(key, (8 // ensemble_size, 16))

    train_loss = train_loss_fn(logits, train_labels)
    eval_loss = eval_loss_fn(logits, eval_labels)

    self.assertEqual(train_loss.shape, train_expected_shape)
    self.assertEqual(eval_loss.shape, eval_expected_shape)
    self.assertEqual(label_pred_fn(logits).shape, eval_expected_shape)
    # When ensemble_size is 1, we are back in the standard case where the train
    # and eval losses are identical.
    if ensemble_size == 1:
      self.assertSequenceEqual(list(train_loss), list(eval_loss))


class GetTrainStepsAndEpochsTest(parameterized.TestCase):

  @parameterized.named_parameters(
      ('_train_epochs', {'train_epochs': 1}, (10, 1.0)),
      ('_train_steps', {'train_steps': 11}, (11, 1.1)),
  )
  def test(self, kwargs, expected):
    self.assertEqual(

        trainer.get_train_steps_and_epochs(