# Path: morph_net/framework/op_handler_util.py
# OP_TYPES_WITH_MULTIPLE_OUTPUTS = ('SplitV',)
# WEIGHTS_INDEX_DICT = {
#     'Conv2D': 1,
#     'Conv2DBackpropInput': 1,
#     'DepthwiseConv2dNative': 1,
#     'MatMul': 1
# }
# def get_input_ops(op, op_reg_manager, whitelist_indices=None):
#   def is_const(tensor):
#   def is_weight_tensor(i, op_type):
# def get_output_ops(op, op_reg_manager):
# def get_ops_without_groups(ops, op_reg_manager):
# def remove_non_passthrough_ops(ops, op_reg_manager):
# def group_op_with_inputs_and_outputs(op, input_op_slices, output_op_slices,
#                                      aligned_op_slice_sizes, op_reg_manager):
# def get_concat_input_op_slices(concat_ops, op_reg_manager):
# def get_op_slices(ops, op_reg_manager):
# def get_op_slice_sizes(op_slices):
# def get_aligned_op_slice_sizes(op_slices, input_op_slices, output_op_slices):
# def get_aligned_sizes(op_slice_sizes):
# def _get_num_slices(op_slice_sizes):
# def reslice_concat_ops(concat_ops, aligned_op_slice_sizes, op_reg_manager):
# def get_total_slice_size(op_slice_sizes, index, slice_count):
# def reslice_ops(ops, aligned_op_slice_sizes, op_reg_manager):
# def _get_source_op_slices(op_slices, op_reg_manager):
# def group_aligned_input_output_slices(
#     op, input_ops_to_process, output_ops_to_process, input_op_slices,
#     output_op_slices, aligned_op_slice_sizes, op_reg_manager):
# def get_op_size(op):
# def separate_same_size_ops(reference_op, ops):
# def group_match(regex, op_slices):
# 
# Path: morph_net/network_regularizers/cost_calculator.py
# CONV2D_OPS = ('Conv2D', 'Conv2DBackpropInput', 'DepthwiseConv2dNative')
# CONV3D_OPS = ('Conv3D',)
# CONV_OPS = CONV2D_OPS + CONV3D_OPS
# FLOP_OPS = CONV_OPS + ('MatMul',)
# SUPPORTED_OPS = FLOP_OPS + ('Add', 'AddV2', 'AddN', 'ConcatV2',
#                             'FusedBatchNorm', 'FusedBatchNormV2',
#                             'FusedBatchNormV3', 'Mul', 'Relu', 'Relu6', 'Sum')
# class CostCalculator(object):
#   def __init__(self, op_regularizer_manager, resource_function):
#   def _get_cost_or_regularization_term(self, is_regularization, ops=None):
#   def get_cost(self, ops=None):
#   def get_regularization_term(self, ops=None):
# def get_input_activation(op):
# def _count_alive(tensor, opreg):
# def _sum_of_reg_vector(opreg):


from morph_net.framework import op_handler_util
from morph_net.network_regularizers import cost_calculator
import numpy as np
import tensorflow.compat.v1 as tf
  del batch_size  # Unused.
  if not is_flop_op(op):
    return 0.0
  if is_regularization:
    return reg_outputs
  return num_alive_outputs


def _shape_with_dtype(tensor):
  """Returns the tensor shape with the same dtype as the tensor.

  Args:
    tensor: A tf.Tensor.

  Returns:
    A tf.Tensor of the tensor shape with the same dtype as tensor.
  """
  return tf.cast(tf.shape(tensor), tensor.dtype)


def is_flop_op(op):
  """Returns True if op consumes significant FLOPs to evaluate."""
  if not isinstance(op, tf.Operation):
    raise ValueError('conv_op must be a tf.Operation, not %s' % type(op))
  return op.type in cost_calculator.FLOP_OPS


def _get_conv_filter_size(conv_op):
  # Works for 2D and 3D convs where sizes of weight matrix are:
  # 4D or 5D tensors: [kernel_size[:], inputs, outputs]

# Path: morph_net/framework/op_handler_util.py
# OP_TYPES_WITH_MULTIPLE_OUTPUTS = ('SplitV',)
# WEIGHTS_INDEX_DICT = {
#     'Conv2D': 1,
#     'Conv2DBackpropInput': 1,
#     'DepthwiseConv2dNative': 1,
#     'MatMul': 1
# }
# def get_input_ops(op, op_reg_manager, whitelist_indices=None):
#   def is_const(tensor):
#   def is_weight_tensor(i, op_type):
# def get_output_ops(op, op_reg_manager):
# def get_ops_without_groups(ops, op_reg_manager):
# def remove_non_passthrough_ops(ops, op_reg_manager):
# def group_op_with_inputs_and_outputs(op, input_op_slices, output_op_slices,
#                                      aligned_op_slice_sizes, op_reg_manager):
# def get_concat_input_op_slices(concat_ops, op_reg_manager):
# def get_op_slices(ops, op_reg_manager):
# def get_op_slice_sizes(op_slices):
# def get_aligned_op_slice_sizes(op_slices, input_op_slices, output_op_slices):
# def get_aligned_sizes(op_slice_sizes):
# def _get_num_slices(op_slice_sizes):
# def reslice_concat_ops(concat_ops, aligned_op_slice_sizes, op_reg_manager):
# def get_total_slice_size(op_slice_sizes, index, slice_count):
# def reslice_ops(ops, aligned_op_slice_sizes, op_reg_manager):
# def _get_source_op_slices(op_slices, op_reg_manager):
# def group_aligned_input_output_slices(
#     op, input_ops_to_process, output_ops_to_process, input_op_slices,
#     output_op_slices, aligned_op_slice_sizes, op_reg_manager):
# def get_op_size(op):
# def separate_same_size_ops(reference_op, ops):
# def group_match(regex, op_slices):
# 
# Path: morph_net/network_regularizers/cost_calculator.py
# CONV2D_OPS = ('Conv2D', 'Conv2DBackpropInput', 'DepthwiseConv2dNative')
# CONV3D_OPS = ('Conv3D',)
# CONV_OPS = CONV2D_OPS + CONV3D_OPS
# FLOP_OPS = CONV_OPS + ('MatMul',)
# SUPPORTED_OPS = FLOP_OPS + ('Add', 'AddV2', 'AddN', 'ConcatV2',
#                             'FusedBatchNorm', 'FusedBatchNormV2',
#                             'FusedBatchNormV3', 'Mul', 'Relu', 'Relu6', 'Sum')
# class CostCalculator(object):
#   def __init__(self, op_regularizer_manager, resource_function):
#   def _get_cost_or_regularization_term(self, is_regularization, ops=None):
#   def get_cost(self, ops=None):
#   def get_regularization_term(self, ops=None):
# def get_input_activation(op):
# def _count_alive(tensor, opreg):
# def _sum_of_reg_vector(opreg):


# Path: morph_net/network_regularizers/resource_function.py
from morph_net.framework import op_handler_util
from morph_net.network_regularizers import cost_calculator
import numpy as np
import tensorflow.compat.v1 as tf

  del batch_size  # Unused.
  if not is_flop_op(op):
    return 0.0
  if is_regularization:
    return reg_outputs
  return num_alive_outputs


def _shape_with_dtype(tensor):
  """Returns the tensor shape with the same dtype as the tensor.

  Args:
    tensor: A tf.Tensor.

  Returns:
    A tf.Tensor of the tensor shape with the same dtype as tensor.
  """
  return tf.cast(tf.shape(tensor), tensor.dtype)


def is_flop_op(op):
  """Returns True if op consumes significant FLOPs to evaluate."""
  if not isinstance(op, tf.Operation):
    raise ValueError('conv_op must be a tf.Operation, not %s' % type(op))
  return op.type in cost_calculator.FLOP_OPS


def _get_conv_filter_size(conv_op):
  # Works for 2D and 3D convs where sizes of weight matrix are:
  # 4D or 5D tensors: [kernel_size[:], inputs, outputs]

  assert conv_op.type in cost_calculator.CONV_OPS