# Path: rigl/experimental/jax/datasets/dataset_base.py
# class Dataset(metaclass=abc.ABCMeta):
# class ImageDataset(Dataset):
#   DATAKEY: Optional[str] = None
#   LABELKEY: str = 'label'
#   DATAKEY = 'image'
#   def __init__(self,
#                name,
#                batch_size,
#                batch_size_test,
#                shuffle_buffer_size,
#                prefetch_size = 1,
#                seed = None):  # pytype: disable=annotation-type-mismatch
#   def _dataset_dir(self):
#   def get_train(self):
#   def get_train_len(self):
#   def get_test(self):
#   def get_test_len(self):
#   def preprocess(
#       self, data):
#   def augment(
#       self, data):
#   def preprocess(
#       self, data):
# 
# 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/pruning/masked.py
# def masked(module, mask):
#   """Convenience function for masking a FLAX module with MaskedModule."""
#   return MaskedModule.partial(wrapped_module=module, mask=mask)
# 
# Path: rigl/experimental/jax/pruning/pruning.py
# def weight_magnitude(weights):
# def prune(
#     model,
#     pruning_rate,
#     saliency_fn = weight_magnitude,
#     mask = None,
#     compare_fn = jnp.greater):
# 
# Path: rigl/experimental/jax/pruning/symmetry.py
# def count_permutations_mask_layer(
#     mask_layer,
#     next_mask_layer = None,
#     parameter_key = 'kernel'):
# def count_permutations_mask(mask):
# def get_mask_stats(mask):
# 
# Path: rigl/experimental/jax/utils/utils.py
# def cross_entropy_loss(log_softmax_logits,
#                        labels):
# def compute_metrics(logits,
#                     labels):
# def _np_converter(obj):
# def dump_dict_json(data_dict, path):
# def count_param(model,
#                 param_names):
# def cosine_similarity(a, b):
# def param_as_array(params):
# def cosine_similarity_model(initial_model,
#                             current_model):
# def vector_difference_norm_model(initial_model,
#                                  current_model):
# def pairwise_longest(iterable):
# T = TypeVar('T')


import collections
import functools
import time
import flax
import jax
import jax.numpy as jnp
import tensorflow.compat.v2 as tf
from typing import Callable, Dict, Mapping, Optional, Tuple, Union
from absl import logging
from flax import jax_utils
from flax.training import common_utils
from rigl.experimental.jax.datasets import dataset_base
from rigl.experimental.jax.models import model_factory
from rigl.experimental.jax.pruning import masked
from rigl.experimental.jax.pruning import pruning
from rigl.experimental.jax.pruning import symmetry
from rigl.experimental.jax.utils import utils
      rng: Random number generator, i.e. jax.random.PRNGKey, to use for model
        training, e.g. dropout.
      summary_writer: An optional tensorboard summary writer for logging
    self._rng = rng

    if self._rng is None:
      self._rng = jax.random.PRNGKey(42)

  def _update_optimizer(self, model: flax.deprecated.nn.Model):
    """Updates the optimizer based on the given model."""
    self.optimizer = jax_utils.replicate(
        self._optimizer_def.create(model))

  def train(
      self,
      num_epochs: int,
      lr_fn: Optional[Callable[[int], float]] = None,
      pruning_rate_fn: Optional[PruningRateFnType] = None,
      update_iter: int = 100,
      update_epoch: int = 10
  ) -> Tuple[flax.deprecated.nn.Model, Mapping[str, Union[int, float, Mapping[
      str, float]]]]:
    """Trains the model over the given number of epochs.

    Args:
      num_epochs: The total number of epochs to train over.
      lr_fn: The learning rate function, takes the current iteration/step as an
        argument and returns the current learning rate, uses constant learning
        rate if no function is provided.
      pruning_rate_fn: The pruning rate function, takes the current epoch as an

# Path: rigl/experimental/jax/datasets/dataset_base.py
# class Dataset(metaclass=abc.ABCMeta):
# class ImageDataset(Dataset):
#   DATAKEY: Optional[str] = None
#   LABELKEY: str = 'label'
#   DATAKEY = 'image'
#   def __init__(self,
#                name,
#                batch_size,
#                batch_size_test,
#                shuffle_buffer_size,
#                prefetch_size = 1,
#                seed = None):  # pytype: disable=annotation-type-mismatch
#   def _dataset_dir(self):
#   def get_train(self):
#   def get_train_len(self):
#   def get_test(self):
#   def get_test_len(self):
#   def preprocess(
#       self, data):
#   def augment(
#       self, data):
#   def preprocess(
#       self, data):
# 
# 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/pruning/masked.py
# def masked(module, mask):
#   """Convenience function for masking a FLAX module with MaskedModule."""
#   return MaskedModule.partial(wrapped_module=module, mask=mask)
# 
# Path: rigl/experimental/jax/pruning/pruning.py
# def weight_magnitude(weights):
# def prune(
#     model,
#     pruning_rate,
#     saliency_fn = weight_magnitude,
#     mask = None,
#     compare_fn = jnp.greater):
# 
# Path: rigl/experimental/jax/pruning/symmetry.py
# def count_permutations_mask_layer(
#     mask_layer,
#     next_mask_layer = None,
#     parameter_key = 'kernel'):
# def count_permutations_mask(mask):
# def get_mask_stats(mask):
# 
# Path: rigl/experimental/jax/utils/utils.py
# def cross_entropy_loss(log_softmax_logits,
#                        labels):
# def compute_metrics(logits,
#                     labels):
# def _np_converter(obj):
# def dump_dict_json(data_dict, path):
# def count_param(model,
#                 param_names):
# def cosine_similarity(a, b):
# def param_as_array(params):
# def cosine_similarity_model(initial_model,
#                             current_model):
# def vector_difference_norm_model(initial_model,
#                                  current_model):
# def pairwise_longest(iterable):
# T = TypeVar('T')


# Path: rigl/experimental/jax/training/training.py
import collections
import functools
import time
import flax
import jax
import jax.numpy as jnp
import tensorflow.compat.v2 as tf
from typing import Callable, Dict, Mapping, Optional, Tuple, Union
from absl import logging
from flax import jax_utils
from flax.training import common_utils
from rigl.experimental.jax.datasets import dataset_base
from rigl.experimental.jax.models import model_factory
from rigl.experimental.jax.pruning import masked
from rigl.experimental.jax.pruning import pruning
from rigl.experimental.jax.pruning import symmetry
from rigl.experimental.jax.utils import utils

      rng: Random number generator, i.e. jax.random.PRNGKey, to use for model
        training, e.g. dropout.
      summary_writer: An optional tensorboard summary writer for logging
    self._rng = rng

    if self._rng is None:
      self._rng = jax.random.PRNGKey(42)

  def _update_optimizer(self, model: flax.deprecated.nn.Model):
    """Updates the optimizer based on the given model."""
    self.optimizer = jax_utils.replicate(
        self._optimizer_def.create(model))

  def train(
      self,
      num_epochs: int,
      lr_fn: Optional[Callable[[int], float]] = None,
      pruning_rate_fn: Optional[PruningRateFnType] = None,
      update_iter: int = 100,
      update_epoch: int = 10
  ) -> Tuple[flax.deprecated.nn.Model, Mapping[str, Union[int, float, Mapping[
      str, float]]]]:
    """Trains the model over the given number of epochs.

    Args:
      num_epochs: The total number of epochs to train over.
      lr_fn: The learning rate function, takes the current iteration/step as an
        argument and returns the current learning rate, uses constant learning
        rate if no function is provided.
      pruning_rate_fn: The pruning rate function, takes the current epoch as an

        argument and returns the current pruning rate, no further pruning is