# 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/pruning/mask_factory.py
# MASK_TYPES: Mapping[str, MaskFnType] = {
#     'random':
#         masked.shuffled_mask,
#     'per_neuron':
#         masked.shuffled_neuron_mask,
#     'per_neuron_no_input_ablation':
#         masked.shuffled_neuron_no_input_ablation_mask,
#     'symmetric':
#         masked.symmetric_mask,
# }
# def create_mask(mask_type, base_model,
#                 rng, sparsity,
#                 **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 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 model_factory
from rigl.experimental.jax.pruning import mask_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
  logging.info('Generating random mask based on model')

  # Re-initialize the RNG to maintain same training pattern (as in prune code).
  mask_rng = jax.random.PRNGKey(FLAGS.mask_randomseed)

  mask = mask_factory.create_mask(FLAGS.mask_type, base_model, mask_rng,
                                  FLAGS.mask_sparsity)

  if jax.host_id() == 0:
    mask_stats = symmetry.get_mask_stats(mask)
    logging.info('Mask stats: %s', str(mask_stats))


    for label, value in mask_stats.items():
      try:
        summary_writer.scalar(f'mask/{label}', value, 0)
      # This is needed because permutations (long int) can't be cast to float32.
      except (OverflowError, ValueError):
        summary_writer.text(f'mask/{label}', str(value), 0)
        logging.error('Could not write mask/%s to tensorflow summary as float32'
                      ', writing as string instead.', label)

    if FLAGS.dump_json:
      mask_stats['permutations'] = str(mask_stats['permutations'])
      utils.dump_dict_json(
          mask_stats, path.join(experiment_dir, 'mask_stats.json'))

  mask = masked.propagate_masks(mask)

  if jax.host_id() == 0:

# 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/pruning/mask_factory.py
# MASK_TYPES: Mapping[str, MaskFnType] = {
#     'random':
#         masked.shuffled_mask,
#     'per_neuron':
#         masked.shuffled_neuron_mask,
#     'per_neuron_no_input_ablation':
#         masked.shuffled_neuron_no_input_ablation_mask,
#     'symmetric':
#         masked.symmetric_mask,
# }
# def create_mask(mask_type, base_model,
#                 rng, sparsity,
#                 **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/random_mask.py
import ast
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 model_factory
from rigl.experimental.jax.pruning import mask_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

  logging.info('Generating random mask based on model')

  # Re-initialize the RNG to maintain same training pattern (as in prune code).
  mask_rng = jax.random.PRNGKey(FLAGS.mask_randomseed)

  mask = mask_factory.create_mask(FLAGS.mask_type, base_model, mask_rng,
                                  FLAGS.mask_sparsity)

  if jax.host_id() == 0:
    mask_stats = symmetry.get_mask_stats(mask)
    logging.info('Mask stats: %s', str(mask_stats))


    for label, value in mask_stats.items():
      try:
        summary_writer.scalar(f'mask/{label}', value, 0)
      # This is needed because permutations (long int) can't be cast to float32.
      except (OverflowError, ValueError):
        summary_writer.text(f'mask/{label}', str(value), 0)
        logging.error('Could not write mask/%s to tensorflow summary as float32'
                      ', writing as string instead.', label)

    if FLAGS.dump_json:
      mask_stats['permutations'] = str(mask_stats['permutations'])
      utils.dump_dict_json(
          mask_stats, path.join(experiment_dir, 'mask_stats.json'))

  mask = masked.propagate_masks(mask)

  if jax.host_id() == 0:

    mask_stats = symmetry.get_mask_stats(mask)