# 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):


from absl.testing import parameterized
from morph_net.network_regularizers import resource_function
from tensorflow.contrib import layers as contrib_layers
from tensorflow.python.framework import ops
import numpy as np
import tensorflow.compat.v1 as tf
    # Expected latency loss = max(compute_loss, memory_loss)
    expected_compute_cost = (
        2 * batch_size * 6 * 7 * 7 * 5 * num_alive_inputs * num_alive_outputs
        / peak_compute)
    expected_memory_cost = (
        (batch_size * 11 * 13 * num_alive_inputs
         + 7 * 5 * num_alive_inputs * num_alive_outputs
         + batch_size * 6 * 7 * num_alive_outputs)
        * self.image.dtype.size / memory_bandwidth)
    expected_compute_loss = (
        2 * batch_size * 6 * 7 * 7 * 5 * (
            num_alive_inputs * reg_outputs + num_alive_outputs * reg_inputs)
        / peak_compute)
    expected_memory_loss = (
        (batch_size * 11 * 13 * reg_inputs
         + 7 * 5 * (
             num_alive_inputs * reg_outputs + num_alive_outputs * reg_inputs)
         + batch_size * 6 * 7 * reg_outputs)
        * self.image.dtype.size / memory_bandwidth)
    if expected_memory_cost > expected_compute_cost:
      expected_latency_loss = expected_memory_loss
    else:
      expected_latency_loss = expected_compute_loss
    self.assertNearRelatively(expected_latency_loss, latency_loss)

  @parameterized.named_parameters(
      ('_BatchSize1_AliveIn10_AliveOut10', 1, 10, 10),
      ('_BatchSize32_AliveIn4_AliveOut9', 32, 4, 9))
  def testDepthwiseConvModelSizeFunction_Cost(
      self, batch_size, num_alive_inputs, num_alive_outputs):

# 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/network_regularizers/resource_function_test.py
from absl.testing import parameterized
from morph_net.network_regularizers import resource_function
from tensorflow.contrib import layers as contrib_layers
from tensorflow.python.framework import ops
import numpy as np
import tensorflow.compat.v1 as tf

    # Expected latency loss = max(compute_loss, memory_loss)
    expected_compute_cost = (
        2 * batch_size * 6 * 7 * 7 * 5 * num_alive_inputs * num_alive_outputs
        / peak_compute)
    expected_memory_cost = (
        (batch_size * 11 * 13 * num_alive_inputs
         + 7 * 5 * num_alive_inputs * num_alive_outputs
         + batch_size * 6 * 7 * num_alive_outputs)
        * self.image.dtype.size / memory_bandwidth)
    expected_compute_loss = (
        2 * batch_size * 6 * 7 * 7 * 5 * (
            num_alive_inputs * reg_outputs + num_alive_outputs * reg_inputs)
        / peak_compute)
    expected_memory_loss = (
        (batch_size * 11 * 13 * reg_inputs
         + 7 * 5 * (
             num_alive_inputs * reg_outputs + num_alive_outputs * reg_inputs)
         + batch_size * 6 * 7 * reg_outputs)
        * self.image.dtype.size / memory_bandwidth)
    if expected_memory_cost > expected_compute_cost:
      expected_latency_loss = expected_memory_loss
    else:
      expected_latency_loss = expected_compute_loss
    self.assertNearRelatively(expected_latency_loss, latency_loss)

  @parameterized.named_parameters(
      ('_BatchSize1_AliveIn10_AliveOut10', 1, 10, 10),
      ('_BatchSize32_AliveIn4_AliveOut9', 32, 4, 9))
  def testDepthwiseConvModelSizeFunction_Cost(
      self, batch_size, num_alive_inputs, num_alive_outputs):

    model_size_cost = resource_function.model_size_function(