# Path: rigl/experimental/jax/datasets/dataset_factory.py
# DATASETS: Mapping[str, Type[dataset_base.Dataset]] = {
#     'MNIST': mnist.MNISTDataset,
#     'CIFAR10': cifar10.CIFAR10Dataset,
# }
# def create_dataset(name, *args, **kwargs):
# 
# Path: rigl/experimental/jax/models/model_factory.py
# MODELS: Mapping[str, Type[flax.deprecated.nn.Model]] = {
#     'MNIST_CNN': mnist_cnn.MNISTCNN,
#     'MNIST_FC': mnist_fc.MNISTFC,
#     'CIFAR10_CNN': cifar10_cnn.CIFAR10CNN,
# }
# def create_model(
#     name, rng,
#     input_specs, **kwargs
# ):
# def update_model(model,
#                  **kwargs):
# 
# Path: rigl/experimental/jax/training/training.py
# LABELKEY = dataset_base.ImageDataset.LABELKEY
# DATAKEY = dataset_base.ImageDataset.DATAKEY
# def _shard_batch(xs):
#   def _prepare(x):
# def train_step(
#     optimizer: flax.optim.Optimizer, batch: Mapping[str, jnp.array],
#     rng: Callable[[int], jnp.array], state: flax.deprecated.nn.Collection,
#     learning_rate_fn: Callable[[int], float]
# ) -> Tuple[flax.optim.Optimizer, flax.deprecated.nn.Collection, float, float]:
#   def loss_fn(
#       model: flax.deprecated.nn.Model
#   ) -> Tuple[float, Tuple[flax.deprecated.nn.Collection, jnp.array]]:
#   def __init__(
#       self,
#       optimizer_def: flax.optim.OptimizerDef,
#       initial_model: flax.deprecated.nn.Model,
#       initial_state: flax.deprecated.nn.Collection,
#       dataset: jnp.array,
#       rng: Callable[[int], jnp.array] = None,
#       summary_writer: Optional[tf.summary.SummaryWriter] = None,
#   ):
# class Trainer:


import functools
import math
import flax
import jax
import jax.numpy as jnp
from absl.testing import absltest
from flax import jax_utils
from flax.metrics import tensorboard
from flax.training import common_utils
from rigl.experimental.jax.datasets import dataset_factory
from rigl.experimental.jax.models import model_factory
from rigl.experimental.jax.training import training
      self.assertBetween(best_metrics['train_avg_loss'], self._min_loss,
                         self._max_loss)

    with self.subTest(name='step'):
      self.assertGreater(best_metrics['step'], 0)

    with self.subTest(name='cosine_distance'):
      self.assertBetween(best_metrics['cosine_distance'], 0., 1.)

    with self.subTest(name='cumulative_gradient_norm'):
      self.assertGreater(best_metrics['cumulative_gradient_norm'], 0)

    with self.subTest(name='test_accuracy'):
      self.assertBetween(best_metrics['test_accuracy'], 0., 1.)

    with self.subTest(name='test_avg_loss'):
      self.assertBetween(best_metrics['test_avg_loss'], self._min_loss,
                         self._max_loss)

  def test_train_one_epoch_tensorboard(self):
    """Tests training loop over one epoch, with tensorboard."""

    trainer = training.Trainer(
        self._optimizer,
        self._model,
        self._state,
        self._dataset,
        summary_writer=self._summarywriter)

    with self.subTest(name='TrainerInstantiation'):

# Path: rigl/experimental/jax/datasets/dataset_factory.py
# DATASETS: Mapping[str, Type[dataset_base.Dataset]] = {
#     'MNIST': mnist.MNISTDataset,
#     'CIFAR10': cifar10.CIFAR10Dataset,
# }
# def create_dataset(name, *args, **kwargs):
# 
# Path: rigl/experimental/jax/models/model_factory.py
# MODELS: Mapping[str, Type[flax.deprecated.nn.Model]] = {
#     'MNIST_CNN': mnist_cnn.MNISTCNN,
#     'MNIST_FC': mnist_fc.MNISTFC,
#     'CIFAR10_CNN': cifar10_cnn.CIFAR10CNN,
# }
# def create_model(
#     name, rng,
#     input_specs, **kwargs
# ):
# def update_model(model,
#                  **kwargs):
# 
# Path: rigl/experimental/jax/training/training.py
# LABELKEY = dataset_base.ImageDataset.LABELKEY
# DATAKEY = dataset_base.ImageDataset.DATAKEY
# def _shard_batch(xs):
#   def _prepare(x):
# def train_step(
#     optimizer: flax.optim.Optimizer, batch: Mapping[str, jnp.array],
#     rng: Callable[[int], jnp.array], state: flax.deprecated.nn.Collection,
#     learning_rate_fn: Callable[[int], float]
# ) -> Tuple[flax.optim.Optimizer, flax.deprecated.nn.Collection, float, float]:
#   def loss_fn(
#       model: flax.deprecated.nn.Model
#   ) -> Tuple[float, Tuple[flax.deprecated.nn.Collection, jnp.array]]:
#   def __init__(
#       self,
#       optimizer_def: flax.optim.OptimizerDef,
#       initial_model: flax.deprecated.nn.Model,
#       initial_state: flax.deprecated.nn.Collection,
#       dataset: jnp.array,
#       rng: Callable[[int], jnp.array] = None,
#       summary_writer: Optional[tf.summary.SummaryWriter] = None,
#   ):
# class Trainer:


# Path: rigl/experimental/jax/training/training_test.py
import functools
import math
import flax
import jax
import jax.numpy as jnp
from absl.testing import absltest
from flax import jax_utils
from flax.metrics import tensorboard
from flax.training import common_utils
from rigl.experimental.jax.datasets import dataset_factory
from rigl.experimental.jax.models import model_factory
from rigl.experimental.jax.training import training

      self.assertBetween(best_metrics['train_avg_loss'], self._min_loss,
                         self._max_loss)

    with self.subTest(name='step'):
      self.assertGreater(best_metrics['step'], 0)

    with self.subTest(name='cosine_distance'):
      self.assertBetween(best_metrics['cosine_distance'], 0., 1.)

    with self.subTest(name='cumulative_gradient_norm'):
      self.assertGreater(best_metrics['cumulative_gradient_norm'], 0)

    with self.subTest(name='test_accuracy'):
      self.assertBetween(best_metrics['test_accuracy'], 0., 1.)

    with self.subTest(name='test_avg_loss'):
      self.assertBetween(best_metrics['test_avg_loss'], self._min_loss,
                         self._max_loss)

  def test_train_one_epoch_tensorboard(self):
    """Tests training loop over one epoch, with tensorboard."""

    trainer = training.Trainer(
        self._optimizer,
        self._model,
        self._state,
        self._dataset,
        summary_writer=self._summarywriter)

    with self.subTest(name='TrainerInstantiation'):

      self.assertIsInstance(trainer, training.Trainer)