# 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/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 collections
import functools
import uuid
import flax
import jax
import jax.numpy as jnp
from os import path
from typing import List
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 model_factory
from rigl.experimental.jax.training import training
from rigl.experimental.jax.utils import utils
  else:
    raise ValueError(f'Unknown LR schedule type {FLAGS.lr_schedule}')

  # Reuses the FLAX learning rate schedule framework for pruning rate schedule.
  pruning_fn_p = functools.partial(
      lr_schedule.create_stepped_learning_rate_schedule, FLAGS.pruning_rate,
      steps_per_epoch)
  if FLAGS.pruning_schedule:
    pruning_schedule = ast.literal_eval(FLAGS.pruning_schedule)
    if isinstance(pruning_schedule, collections.Mapping):
      pruning_rate_fn = {
          f'MaskedModule_{layer_num}': pruning_fn_p(schedule)
          for layer_num, schedule in pruning_schedule.items()
      }
    else:
      pruning_rate_fn = pruning_fn_p(pruning_schedule)
  else:
    pruning_rate_fn = lr_schedule.create_constant_learning_rate_schedule(
        FLAGS.pruning_rate, steps_per_epoch)

  if jax.host_id() == 0:
    trainer = training.Trainer(
        optimizer,
        initial_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/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/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/prune.py
import ast
import collections
import functools
import uuid
import flax
import jax
import jax.numpy as jnp
from os import path
from typing import List
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 model_factory
from rigl.experimental.jax.training import training
from rigl.experimental.jax.utils import utils

  else:
    raise ValueError(f'Unknown LR schedule type {FLAGS.lr_schedule}')

  # Reuses the FLAX learning rate schedule framework for pruning rate schedule.
  pruning_fn_p = functools.partial(
      lr_schedule.create_stepped_learning_rate_schedule, FLAGS.pruning_rate,
      steps_per_epoch)
  if FLAGS.pruning_schedule:
    pruning_schedule = ast.literal_eval(FLAGS.pruning_schedule)
    if isinstance(pruning_schedule, collections.Mapping):
      pruning_rate_fn = {
          f'MaskedModule_{layer_num}': pruning_fn_p(schedule)
          for layer_num, schedule in pruning_schedule.items()
      }
    else:
      pruning_rate_fn = pruning_fn_p(pruning_schedule)
  else:
    pruning_rate_fn = lr_schedule.create_constant_learning_rate_schedule(
        FLAGS.pruning_rate, steps_per_epoch)

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

    trainer = training.Trainer(