# 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._dataset_name = 'MNIST'
    self._model_name = 'MNIST_CNN'

    self._summarywriter = tensorboard.SummaryWriter('/tmp/')

    self._dataset = dataset_factory.create_dataset(
        self._dataset_name,
        self._batch_size,
        self._batch_size_test,
        shuffle_buffer_size=self._shuffle_buffer_size)

    self._model, self._state = model_factory.create_model(
        self._model_name,
        self._rng, (self._input_shape,),
        num_classes=self._num_classes)

    self._optimizer = flax.optim.Momentum(
        learning_rate=self._learning_rate_fn(0),
        beta=self._momentum,
        weight_decay=self._weight_decay)

  def test_train_one_step(self):
    """Tests training loop over one step."""
    iterator = self._dataset.get_train()
    batch = next(iterator)

    state = jax_utils.replicate(self._state)
    optimizer = jax_utils.replicate(self._optimizer.create(self._model))

    self._rng, step_key = jax.random.split(self._rng)

# 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._dataset_name = 'MNIST'
    self._model_name = 'MNIST_CNN'

    self._summarywriter = tensorboard.SummaryWriter('/tmp/')

    self._dataset = dataset_factory.create_dataset(
        self._dataset_name,
        self._batch_size,
        self._batch_size_test,
        shuffle_buffer_size=self._shuffle_buffer_size)

    self._model, self._state = model_factory.create_model(
        self._model_name,
        self._rng, (self._input_shape,),
        num_classes=self._num_classes)

    self._optimizer = flax.optim.Momentum(
        learning_rate=self._learning_rate_fn(0),
        beta=self._momentum,
        weight_decay=self._weight_decay)

  def test_train_one_step(self):
    """Tests training loop over one step."""
    iterator = self._dataset.get_train()
    batch = next(iterator)

    state = jax_utils.replicate(self._state)
    optimizer = jax_utils.replicate(self._optimizer.create(self._model))

    self._rng, step_key = jax.random.split(self._rng)

    batch = training._shard_batch(batch)