# 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 gin
import tensorflow as tf
from rigl.rigl_tf2 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.

"""Implements RigL."""


def get_all_layers(model, filter_fn=lambda _: True):
  """Gets all layers of a model and layers of a layer if it is a keras.Model."""
  all_layers = []
  for l in model.layers:
    if hasattr(l, 'layers'):
      all_layers.extend(get_all_layers(l, filter_fn=filter_fn))
    elif filter_fn(l):
      all_layers.append(l)
  return all_layers


def is_pruned(layer):

# 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/mask_updaters.py
import gin
import tensorflow as tf
from rigl.rigl_tf2 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.

"""Implements RigL."""


def get_all_layers(model, filter_fn=lambda _: True):
  """Gets all layers of a model and layers of a layer if it is a keras.Model."""
  all_layers = []
  for l in model.layers:
    if hasattr(l, 'layers'):
      all_layers.extend(get_all_layers(l, filter_fn=filter_fn))
    elif filter_fn(l):
      all_layers.append(l)
  return all_layers


def is_pruned(layer):

  return isinstance(layer, utils.PRUNING_WRAPPER) and layer.trainable