# Path: morph_net/network_regularizers/flop_regularizer.py
# class LogisticSigmoidFlopsRegularizer(
#     logistic_sigmoid_regularizer.LogisticSigmoidRegularizer):
# class GammaFlopsRegularizer(generic_regularizers.NetworkRegularizer):
# class GroupLassoFlopsRegularizer(generic_regularizers.NetworkRegularizer):
#   def get_calculator(self):
#   def name(self):
#   def cost_name(self):
#   def __init__(self,
#                output_boundary: List[tf.Operation],
#                gamma_threshold,
#                regularizer_decorator: Optional[Type[
#                    generic_regularizers.OpRegularizer]] = None,
#                decorator_parameters=None,
#                input_boundary: Optional[List[tf.Operation]] = None,
#                force_group=None,
#                regularizer_blacklist=None):
#   def get_regularization_term(self, ops=None):
#   def get_cost(self, ops=None):
#   def op_regularizer_manager(self):
#   def name(self):
#   def cost_name(self):
#   def __init__(self,
#                output_boundary: List[tf.Operation],
#                threshold,
#                l1_fraction=0,
#                regularizer_decorator: Optional[Type[
#                    generic_regularizers.OpRegularizer]] = None,
#                decorator_parameters=None,
#                input_boundary: Optional[List[tf.Operation]] = None,
#                force_group=None,
#                regularizer_blacklist=None):
#   def get_regularization_term(self, ops=None):
#   def get_cost(self, ops=None):
#   def op_regularizer_manager(self):
#   def name(self):
#   def cost_name(self):
# 
# Path: morph_net/network_regularizers/resource_function.py
# PEAK_COMPUTE = {  # GFLOP/s
#     'K80': 8740,
#     'P4': 5500,
#     'T4': 8100,
#     'P100': 9300,
#     'V100': 125000,
#     'TPUv2': 22500,
#     'FLOP_LATENCY': 1,  # Simulate a device with infinite memory bandwidth.
#     'MEMORY_LATENCY': 1e20,  # Simulate a device with infinite peak compute.
# }
# MEMORY_BANDWIDTH = {  # GB/s
#     'K80': 480,
#     'P4': 192,
#     'T4': 300,
#     'P100': 732,
#     'V100': 900,
#     'TPUv2': 300,
#     'FLOP_LATENCY': 1e20,
#     'MEMORY_LATENCY': 1,
# }
# def flop_coeff(op):
#   def kernel_num_elements(tensor):
# def num_weights_coeff(op):
# def flop_function(op, is_regularization, num_alive_inputs, num_alive_outputs,
#                   reg_inputs, reg_outputs, batch_size=1):
# def memory_function(op, is_regularization, num_alive_inputs, num_alive_outputs,
#                     reg_inputs, reg_outputs, batch_size=1):
# def latency_function(op, is_regularization, num_alive_inputs, num_alive_outputs,
#                      reg_inputs, reg_outputs, peak_compute, memory_bandwidth,
#                      batch_size=1):
# def latency_function_factory(hardware, batch_size):
#   def latency_function_for_hardware(
#       op, is_regularization, num_alive_inputs, num_alive_outputs, reg_inputs,
#       reg_outputs):
# def model_size_function(op, is_regularization, num_alive_inputs,
#                         num_alive_outputs, reg_inputs, reg_outputs,
#                         batch_size=1):
# def activation_count_function(op, is_regularization, num_alive_inputs,
#                               num_alive_outputs, reg_inputs, reg_outputs,
#                               batch_size=1):
# def _shape_with_dtype(tensor):
# def is_flop_op(op):
# def _get_conv_filter_size(conv_op):
# def _calculate_bilinear_regularization(
#     op, coeff, num_alive_inputs, num_alive_outputs, reg_inputs, reg_outputs,
#     batch_size):
# def _calculate_bilinear_cost(
#     op, coeff, num_alive_inputs, num_alive_outputs, batch_size):
# 
# Path: morph_net/testing/dummy_decorator.py
# class DummyDecorator(generic_regularizers.OpRegularizer):
#   def __init__(self, regularizer_object, scale):
#   def regularization_vector(self):
#   def alive_vector(self):


from absl.testing import parameterized
from morph_net.network_regularizers import flop_regularizer
from morph_net.network_regularizers import resource_function
from morph_net.testing import dummy_decorator
from tensorflow.contrib import layers as contrib_layers
from tensorflow.contrib import slim as contrib_slim
from tensorflow.contrib.slim.nets import resnet_v1
import numpy as np
import tensorflow.compat.v1 as tf

  def testFlopRegularizer(self):
    tf.reset_default_graph()
    tf.set_random_seed(7907)
    with slim.arg_scope(
        [slim.layers.conv2d, slim.layers.conv2d_transpose],
        weights_initializer=tf.random_normal_initializer):
      # Our test model is:
      #
      #         -> conv1 --+
      #        /           |--[concat]
      #  image --> conv2 --+
      #        \
      #         -> convt
      #
      # (the model has two "outputs", convt and concat).
      #
      image = tf.constant(0.0, shape=[1, 17, 19, NUM_CHANNELS])
      conv1 = slim.layers.conv2d(
          image, 13, [7, 5], padding='SAME', scope='conv1')
      conv2 = slim.layers.conv2d(
          image, 23, [1, 1], padding='SAME', scope='conv2')
      self.concat = tf.concat([conv1, conv2], 3)
      self.convt = slim.layers.conv2d_transpose(
          image, 29, [7, 5], stride=3, padding='SAME', scope='convt')
      self.name_to_var = {v.op.name: v for v in tf.global_variables()}
    with self.cached_session():
      tf.global_variables_initializer().run()

    threshold = 1.0

# Path: morph_net/network_regularizers/flop_regularizer.py
# class LogisticSigmoidFlopsRegularizer(
#     logistic_sigmoid_regularizer.LogisticSigmoidRegularizer):
# class GammaFlopsRegularizer(generic_regularizers.NetworkRegularizer):
# class GroupLassoFlopsRegularizer(generic_regularizers.NetworkRegularizer):
#   def get_calculator(self):
#   def name(self):
#   def cost_name(self):
#   def __init__(self,
#                output_boundary: List[tf.Operation],
#                gamma_threshold,
#                regularizer_decorator: Optional[Type[
#                    generic_regularizers.OpRegularizer]] = None,
#                decorator_parameters=None,
#                input_boundary: Optional[List[tf.Operation]] = None,
#                force_group=None,
#                regularizer_blacklist=None):
#   def get_regularization_term(self, ops=None):
#   def get_cost(self, ops=None):
#   def op_regularizer_manager(self):
#   def name(self):
#   def cost_name(self):
#   def __init__(self,
#                output_boundary: List[tf.Operation],
#                threshold,
#                l1_fraction=0,
#                regularizer_decorator: Optional[Type[
#                    generic_regularizers.OpRegularizer]] = None,
#                decorator_parameters=None,
#                input_boundary: Optional[List[tf.Operation]] = None,
#                force_group=None,
#                regularizer_blacklist=None):
#   def get_regularization_term(self, ops=None):
#   def get_cost(self, ops=None):
#   def op_regularizer_manager(self):
#   def name(self):
#   def cost_name(self):
# 
# Path: morph_net/network_regularizers/resource_function.py
# PEAK_COMPUTE = {  # GFLOP/s
#     'K80': 8740,
#     'P4': 5500,
#     'T4': 8100,
#     'P100': 9300,
#     'V100': 125000,
#     'TPUv2': 22500,
#     'FLOP_LATENCY': 1,  # Simulate a device with infinite memory bandwidth.
#     'MEMORY_LATENCY': 1e20,  # Simulate a device with infinite peak compute.
# }
# MEMORY_BANDWIDTH = {  # GB/s
#     'K80': 480,
#     'P4': 192,
#     'T4': 300,
#     'P100': 732,
#     'V100': 900,
#     'TPUv2': 300,
#     'FLOP_LATENCY': 1e20,
#     'MEMORY_LATENCY': 1,
# }
# def flop_coeff(op):
#   def kernel_num_elements(tensor):
# def num_weights_coeff(op):
# def flop_function(op, is_regularization, num_alive_inputs, num_alive_outputs,
#                   reg_inputs, reg_outputs, batch_size=1):
# def memory_function(op, is_regularization, num_alive_inputs, num_alive_outputs,
#                     reg_inputs, reg_outputs, batch_size=1):
# def latency_function(op, is_regularization, num_alive_inputs, num_alive_outputs,
#                      reg_inputs, reg_outputs, peak_compute, memory_bandwidth,
#                      batch_size=1):
# def latency_function_factory(hardware, batch_size):
#   def latency_function_for_hardware(
#       op, is_regularization, num_alive_inputs, num_alive_outputs, reg_inputs,
#       reg_outputs):
# def model_size_function(op, is_regularization, num_alive_inputs,
#                         num_alive_outputs, reg_inputs, reg_outputs,
#                         batch_size=1):
# def activation_count_function(op, is_regularization, num_alive_inputs,
#                               num_alive_outputs, reg_inputs, reg_outputs,
#                               batch_size=1):
# def _shape_with_dtype(tensor):
# def is_flop_op(op):
# def _get_conv_filter_size(conv_op):
# def _calculate_bilinear_regularization(
#     op, coeff, num_alive_inputs, num_alive_outputs, reg_inputs, reg_outputs,
#     batch_size):
# def _calculate_bilinear_cost(
#     op, coeff, num_alive_inputs, num_alive_outputs, batch_size):
# 
# Path: morph_net/testing/dummy_decorator.py
# class DummyDecorator(generic_regularizers.OpRegularizer):
#   def __init__(self, regularizer_object, scale):
#   def regularization_vector(self):
#   def alive_vector(self):


# Path: morph_net/network_regularizers/flop_regularizer_test.py
from absl.testing import parameterized
from morph_net.network_regularizers import flop_regularizer
from morph_net.network_regularizers import resource_function
from morph_net.testing import dummy_decorator
from tensorflow.contrib import layers as contrib_layers
from tensorflow.contrib import slim as contrib_slim
from tensorflow.contrib.slim.nets import resnet_v1
import numpy as np
import tensorflow.compat.v1 as tf


  def testFlopRegularizer(self):
    tf.reset_default_graph()
    tf.set_random_seed(7907)
    with slim.arg_scope(
        [slim.layers.conv2d, slim.layers.conv2d_transpose],
        weights_initializer=tf.random_normal_initializer):
      # Our test model is:
      #
      #         -> conv1 --+
      #        /           |--[concat]
      #  image --> conv2 --+
      #        \
      #         -> convt
      #
      # (the model has two "outputs", convt and concat).
      #
      image = tf.constant(0.0, shape=[1, 17, 19, NUM_CHANNELS])
      conv1 = slim.layers.conv2d(
          image, 13, [7, 5], padding='SAME', scope='conv1')
      conv2 = slim.layers.conv2d(
          image, 23, [1, 1], padding='SAME', scope='conv2')
      self.concat = tf.concat([conv1, conv2], 3)
      self.convt = slim.layers.conv2d_transpose(
          image, 29, [7, 5], stride=3, padding='SAME', scope='convt')
      self.name_to_var = {v.op.name: v for v in tf.global_variables()}
    with self.cached_session():
      tf.global_variables_initializer().run()

    threshold = 1.0

    flop_reg = flop_regularizer.GroupLassoFlopsRegularizer(