# Path: rigl/rigl_tf2/utils.py
# FLAGS = flags.FLAGS
# PRUNING_WRAPPER = pruning_wrapper.PruneLowMagnitude
# PRUNED_LAYER_TYPES = (tf.keras.layers.Conv2D, tf.keras.layers.Dense)
# def get_dataset():
# def get_pruning_params(mode='prune',
#                        initial_sparsity=0.0,
#                        final_sparsity=0.8,
#                        begin_step=2000,
#                        end_step=4000,
#                        frequency=200):
# def maybe_prune_layer(layer, params, filter_fn):
# 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):
# def get_optimizer(total_steps,
#                   name = 'adam',
#                   learning_rate = 0.001,
#                   clipnorm = None,
#                   clipvalue = None,
#                   momentum = None):


import os
import gin
import numpy as np
import tensorflow.compat.v2 as tf
from absl import app
from absl import flags
from absl import logging
from rigl.rigl_tf2 import utils
from pyglib import timer
                                             model_inter.trainable_variables):
      new_value = (1 - i_coef) * var_start + i_coef * var_end
      var_inter.assign(new_value)
    scores = test_model(model_inter, d_set)
    all_scores[i_coef] = scores
  return all_scores


def main(unused_argv):
  init_timer = timer.Timer()
  init_timer.Start()
  if FLAGS.preload_gin_config:
    # Load default values from the original experiment, always the first one.
    with gin.unlock_config():
      gin.parse_config_file(FLAGS.preload_gin_config, skip_unknown=True)
    logging.info('Operative Gin configurations loaded from: %s',
                 FLAGS.preload_gin_config)
  gin.parse_config_files_and_bindings(FLAGS.gin_config, FLAGS.gin_bindings)

  data_train, data_test, info = utils.get_dataset()
  input_shape = info.features['image'].shape
  num_classes = info.features['label'].num_classes
  logging.info('Input Shape: %s', input_shape)
  logging.info('train samples: %s', info.splits['train'].num_examples)
  logging.info('test samples: %s', info.splits['test'].num_examples)
  data_eval = data_train if FLAGS.eval_on_train else data_test
  pruning_params = utils.get_pruning_params(mode='constant')
  mask_load_dict = {-1: None, 0: FLAGS.ckpt_start, 1: FLAGS.ckpt_end}
  mask_path = mask_load_dict[FLAGS.load_mask_from]
  # Currently we interpolate only on the same sparse space.

# Path: rigl/rigl_tf2/utils.py
# FLAGS = flags.FLAGS
# PRUNING_WRAPPER = pruning_wrapper.PruneLowMagnitude
# PRUNED_LAYER_TYPES = (tf.keras.layers.Conv2D, tf.keras.layers.Dense)
# def get_dataset():
# def get_pruning_params(mode='prune',
#                        initial_sparsity=0.0,
#                        final_sparsity=0.8,
#                        begin_step=2000,
#                        end_step=4000,
#                        frequency=200):
# def maybe_prune_layer(layer, params, filter_fn):
# 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):
# def get_optimizer(total_steps,
#                   name = 'adam',
#                   learning_rate = 0.001,
#                   clipnorm = None,
#                   clipvalue = None,
#                   momentum = None):


# Path: rigl/rigl_tf2/interpolate.py
import os
import gin
import numpy as np
import tensorflow.compat.v2 as tf
from absl import app
from absl import flags
from absl import logging
from rigl.rigl_tf2 import utils
from pyglib import timer

                                             model_inter.trainable_variables):
      new_value = (1 - i_coef) * var_start + i_coef * var_end
      var_inter.assign(new_value)
    scores = test_model(model_inter, d_set)
    all_scores[i_coef] = scores
  return all_scores


def main(unused_argv):
  init_timer = timer.Timer()
  init_timer.Start()
  if FLAGS.preload_gin_config:
    # Load default values from the original experiment, always the first one.
    with gin.unlock_config():
      gin.parse_config_file(FLAGS.preload_gin_config, skip_unknown=True)
    logging.info('Operative Gin configurations loaded from: %s',
                 FLAGS.preload_gin_config)
  gin.parse_config_files_and_bindings(FLAGS.gin_config, FLAGS.gin_bindings)

  data_train, data_test, info = utils.get_dataset()
  input_shape = info.features['image'].shape
  num_classes = info.features['label'].num_classes
  logging.info('Input Shape: %s', input_shape)
  logging.info('train samples: %s', info.splits['train'].num_examples)
  logging.info('test samples: %s', info.splits['test'].num_examples)
  data_eval = data_train if FLAGS.eval_on_train else data_test
  pruning_params = utils.get_pruning_params(mode='constant')
  mask_load_dict = {-1: None, 0: FLAGS.ckpt_start, 1: FLAGS.ckpt_end}
  mask_path = mask_load_dict[FLAGS.load_mask_from]
  # Currently we interpolate only on the same sparse space.

  model_start = utils.get_network(