# 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
  else:
    raise ValueError('Mode: %s, is not valid' % mode)
  return p_params


# Forked from tensorflow_model_optimization/python/core/sparsity/keras/prune.py
def maybe_prune_layer(layer, params, filter_fn):
  if filter_fn(layer):
    return PRUNING_WRAPPER(layer, **params)
  return layer


@gin.configurable('network')
def get_network(
    pruning_params,
    input_shape,
    num_classes,
    activation = 'relu',
    network_name = 'lenet5',
    mask_init_path = None,
    shuffle_mask = False,
    weight_init_path = None,
    weight_init_method = None,
    weight_decay = 0.,
    noise_stddev = 0.,
    pruned_layer_types = PRUNED_LAYER_TYPES):
  """Creates the network."""
  kernel_regularizer = (
      tf.keras.regularizers.l2(weight_decay) if (weight_decay > 0) else None)
  # (1) Create keras model.

# 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

  else:
    raise ValueError('Mode: %s, is not valid' % mode)
  return p_params


# Forked from tensorflow_model_optimization/python/core/sparsity/keras/prune.py
def maybe_prune_layer(layer, params, filter_fn):
  if filter_fn(layer):
    return PRUNING_WRAPPER(layer, **params)
  return layer


@gin.configurable('network')
def get_network(
    pruning_params,
    input_shape,
    num_classes,
    activation = 'relu',
    network_name = 'lenet5',
    mask_init_path = None,
    shuffle_mask = False,
    weight_init_path = None,
    weight_init_method = None,
    weight_decay = 0.,
    noise_stddev = 0.,
    pruned_layer_types = PRUNED_LAYER_TYPES):
  """Creates the network."""
  kernel_regularizer = (
      tf.keras.regularizers.l2(weight_decay) if (weight_decay > 0) else None)
  # (1) Create keras model.

  model = getattr(networks, network_name)(