# 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/mnist_fc.py
# def feature_dim_for_param(input_len,
#                           param_count,
#                           depth,
#                           depth_mult = 2.):
#   def apply(self,
#             inputs,
#             num_classes,
#             features = (32, 32),
#             train=True,
#             init_fn = flax.deprecated.nn.initializers.kaiming_normal,
#             activation_fn = flax.deprecated.nn.relu,
#             masks = None,
#             masked_layer_indices = None,
#             dropout_rate = 0.):
# class MNISTFC(flax.deprecated.nn.Module):
# 
# 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/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/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/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 ast
import functools
import operator
import uuid
import flax
import jax
import jax.numpy as jnp
from os import path
from typing import List, Sequence
from absl import app
from absl import flags
from absl import logging
from flax.metrics import tensorboard
from flax.training import lr_schedule
from rigl.experimental.jax.datasets import dataset_factory
from rigl.experimental.jax.models import mnist_fc
from rigl.experimental.jax.models import model_factory
from rigl.experimental.jax.pruning import masked
from rigl.experimental.jax.pruning import symmetry
from rigl.experimental.jax.training import training
from rigl.experimental.jax.utils import utils
        weight_decay=FLAGS.weight_decay,
        nesterov=False)
  else:
    raise ValueError('Unknown Optimizer: {}'.format(FLAGS.opt))

  steps_per_epoch = dataset.get_train_len() // FLAGS.batch_size

  if FLAGS.lr_schedule == 'constant':
    lr_fn = lr_schedule.create_constant_learning_rate_schedule(
        FLAGS.lr, steps_per_epoch)
  elif FLAGS.lr_schedule == 'stepped':
    lr_schedule_steps = ast.literal_eval(FLAGS.lr_schedule_steps)
    lr_fn = lr_schedule.create_stepped_learning_rate_schedule(
        FLAGS.lr, steps_per_epoch, lr_schedule_steps)
  elif FLAGS.lr_schedule == 'cosine':
    lr_fn = lr_schedule.create_cosine_learning_rate_schedule(
        FLAGS.lr, steps_per_epoch, FLAGS.epochs)
  else:
    raise ValueError('Unknown LR schedule type {}'.format(FLAGS.lr_schedule))

  if jax.host_id() == 0:
    trainer = training.Trainer(
        optimizer,
        model,
        initial_state,
        dataset,
        rng,
        summary_writer=summary_writer,
    )
  else:

# 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/mnist_fc.py
# def feature_dim_for_param(input_len,
#                           param_count,
#                           depth,
#                           depth_mult = 2.):
#   def apply(self,
#             inputs,
#             num_classes,
#             features = (32, 32),
#             train=True,
#             init_fn = flax.deprecated.nn.initializers.kaiming_normal,
#             activation_fn = flax.deprecated.nn.relu,
#             masks = None,
#             masked_layer_indices = None,
#             dropout_rate = 0.):
# class MNISTFC(flax.deprecated.nn.Module):
# 
# 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/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/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/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/fixed_param.py
import ast
import functools
import operator
import uuid
import flax
import jax
import jax.numpy as jnp
from os import path
from typing import List, Sequence
from absl import app
from absl import flags
from absl import logging
from flax.metrics import tensorboard
from flax.training import lr_schedule
from rigl.experimental.jax.datasets import dataset_factory
from rigl.experimental.jax.models import mnist_fc
from rigl.experimental.jax.models import model_factory
from rigl.experimental.jax.pruning import masked
from rigl.experimental.jax.pruning import symmetry
from rigl.experimental.jax.training import training
from rigl.experimental.jax.utils import utils

        weight_decay=FLAGS.weight_decay,
        nesterov=False)
  else:
    raise ValueError('Unknown Optimizer: {}'.format(FLAGS.opt))

  steps_per_epoch = dataset.get_train_len() // FLAGS.batch_size

  if FLAGS.lr_schedule == 'constant':
    lr_fn = lr_schedule.create_constant_learning_rate_schedule(
        FLAGS.lr, steps_per_epoch)
  elif FLAGS.lr_schedule == 'stepped':
    lr_schedule_steps = ast.literal_eval(FLAGS.lr_schedule_steps)
    lr_fn = lr_schedule.create_stepped_learning_rate_schedule(
        FLAGS.lr, steps_per_epoch, lr_schedule_steps)
  elif FLAGS.lr_schedule == 'cosine':
    lr_fn = lr_schedule.create_cosine_learning_rate_schedule(
        FLAGS.lr, steps_per_epoch, FLAGS.epochs)
  else:
    raise ValueError('Unknown LR schedule type {}'.format(FLAGS.lr_schedule))

  if jax.host_id() == 0:
    trainer = training.Trainer(
        optimizer,
        model,
        initial_state,
        dataset,
        rng,
        summary_writer=summary_writer,
    )
  else:

    trainer = training.Trainer(optimizer, model, initial_state, dataset, rng)