# Path: rigl/rigl_tf2/init_utils.py
# def unit_scaled_init(mask, method='fanavg_uniform', scale=1.0):
# def layer_scaled_init(mask, method='fanavg_uniform', scale=1.0):
# def unit_scaled_init_tf1(mask,
#                          method='fanavg_uniform',
#                          scale=1.0,
#                          dtype=tf.float32):
#   def new_val_fn(index):
# 
# Path: rigl/rigl_tf2/networks.py
# def lenet5(input_shape,
#            num_classes,
#            activation,
#            kernel_regularizer,
#            use_batch_norm = False,
#            hidden_sizes = (6, 16, 120, 84)):
#   def maybe_add_batchnorm():
# def mlp(input_shape,
#         num_classes,
#         activation,
#         kernel_regularizer,
#         use_batch_norm = False,
#         hidden_sizes = (300, 100)):
#   def maybe_add_batchnorm():


import functools
import gin
import tensorflow.compat.v2 as tf
import tensorflow_datasets as tfds
from typing import Optional, Tuple
from absl import flags
from absl import logging
from rigl.rigl_tf2 import init_utils
from rigl.rigl_tf2 import networks
from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule
from tensorflow_model_optimization.python.core.sparsity.keras import pruning_wrapper
  if weight_init_path:
    logging.info('Loading weights from: %s', weight_init_path)
    weight_init_model = tf.keras.models.clone_model(model)
    ckpt = tf.train.Checkpoint(model=weight_init_model)
    ckpt.restore(weight_init_path)
    for l_source, l_target in zip(weight_init_model.layers, model.layers):
      for var_source, var_target in zip(l_source.trainable_variables,
                                        l_target.trainable_variables):
        var_target.assign(var_source)
        logging.info('Weight %s loaded from ckpt.', var_target.name)
    del weight_init_model
  elif weight_init_method == 'unit_scaled':
    logging.info('Using unit_scaled initialization.')
    for layer in model.layers:
      if isinstance(layer, PRUNING_WRAPPER):
        # TODO following the outcome of b/148083099, update following.
        # Add the weight, mask and the valid dimensions.
        weight = layer.weights[0]
        mask = layer.weights[2]
        new_init = init_utils.unit_scaled_init(mask)
        weight.assign(new_init)
        logging.info('Weight %s updated init.', weight.name)
  elif weight_init_method == 'layer_scaled':
    logging.info('Using layer_scaled initialization.')
    for layer in model.layers:
      if isinstance(layer, PRUNING_WRAPPER):
        # TODO following the outcome of b/148083099, update following.
        # Add the weight, mask and the valid dimensions.
        weight = layer.weights[0]
        mask = layer.weights[2]

# Path: rigl/rigl_tf2/init_utils.py
# def unit_scaled_init(mask, method='fanavg_uniform', scale=1.0):
# def layer_scaled_init(mask, method='fanavg_uniform', scale=1.0):
# def unit_scaled_init_tf1(mask,
#                          method='fanavg_uniform',
#                          scale=1.0,
#                          dtype=tf.float32):
#   def new_val_fn(index):
# 
# Path: rigl/rigl_tf2/networks.py
# def lenet5(input_shape,
#            num_classes,
#            activation,
#            kernel_regularizer,
#            use_batch_norm = False,
#            hidden_sizes = (6, 16, 120, 84)):
#   def maybe_add_batchnorm():
# def mlp(input_shape,
#         num_classes,
#         activation,
#         kernel_regularizer,
#         use_batch_norm = False,
#         hidden_sizes = (300, 100)):
#   def maybe_add_batchnorm():


# Path: rigl/rigl_tf2/utils.py
import functools
import gin
import tensorflow.compat.v2 as tf
import tensorflow_datasets as tfds
from typing import Optional, Tuple
from absl import flags
from absl import logging
from rigl.rigl_tf2 import init_utils
from rigl.rigl_tf2 import networks
from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule
from tensorflow_model_optimization.python.core.sparsity.keras import pruning_wrapper

  if weight_init_path:
    logging.info('Loading weights from: %s', weight_init_path)
    weight_init_model = tf.keras.models.clone_model(model)
    ckpt = tf.train.Checkpoint(model=weight_init_model)
    ckpt.restore(weight_init_path)
    for l_source, l_target in zip(weight_init_model.layers, model.layers):
      for var_source, var_target in zip(l_source.trainable_variables,
                                        l_target.trainable_variables):
        var_target.assign(var_source)
        logging.info('Weight %s loaded from ckpt.', var_target.name)
    del weight_init_model
  elif weight_init_method == 'unit_scaled':
    logging.info('Using unit_scaled initialization.')
    for layer in model.layers:
      if isinstance(layer, PRUNING_WRAPPER):
        # TODO following the outcome of b/148083099, update following.
        # Add the weight, mask and the valid dimensions.
        weight = layer.weights[0]
        mask = layer.weights[2]
        new_init = init_utils.unit_scaled_init(mask)
        weight.assign(new_init)
        logging.info('Weight %s updated init.', weight.name)
  elif weight_init_method == 'layer_scaled':
    logging.info('Using layer_scaled initialization.')
    for layer in model.layers:
      if isinstance(layer, PRUNING_WRAPPER):
        # TODO following the outcome of b/148083099, update following.
        # Add the weight, mask and the valid dimensions.
        weight = layer.weights[0]
        mask = layer.weights[2]

        new_init = init_utils.layer_scaled_init(mask)