# 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
# coding=utf-8
# Copyright 2022 RigL Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Common training code.

This module contains utility functions for training NN.

Attributes:
  LABELKEY: The key used to retrieve a label from the batch dictionary.
  DATAKEY: The key used to retrieve an input image from the batch dictionary.
  PruningRateFnType: Typing alias for a valid pruning rate function.
"""


LABELKEY = dataset_base.ImageDataset.LABELKEY

# 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

# coding=utf-8
# Copyright 2022 RigL Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Common training code.

This module contains utility functions for training NN.

Attributes:
  LABELKEY: The key used to retrieve a label from the batch dictionary.
  DATAKEY: The key used to retrieve an input image from the batch dictionary.
  PruningRateFnType: Typing alias for a valid pruning rate function.
"""


LABELKEY = dataset_base.ImageDataset.LABELKEY

DATAKEY = dataset_base.ImageDataset.DATAKEY