# Path: morph_net/framework/concat_op_handler.py
# def _get_concat_op_axis(op):
#   def is_source_op(self):
#   def is_passthrough(self):
#   def assign_grouping(self, op, op_reg_manager):
#   def create_regularizer(self, _):
#   def _get_input_output_op_slices(self, input_ops, output_ops, op_reg_manager):
# class ConcatOpHandler(op_handler.OpHandler):
# 
# Path: morph_net/framework/depthwise_convolution_op_handler.py
# class DepthwiseConvolutionOpHandler(grouping_op_handler.GroupingOpHandler):
#   def assign_grouping(self, op, op_reg_manager):
#   def _get_depth_multiplier_input_output_op_slices(
#       self, input_ops, input_size, output_ops, op_reg_manager,
#       depth_multiplier):
# 
# Path: morph_net/framework/grouping_op_handler.py
# class GroupingOpHandler(op_handler.OpHandler):
#   def __init__(self, grouping_indices=None):
#   def is_source_op(self):
#   def is_passthrough(self):
#   def assign_grouping(self, op, op_reg_manager):
#   def create_regularizer(self, _):
#   def _get_input_output_op_slices(self, input_ops, output_ops, op_reg_manager):
#   def _is_broadcast(self, op, op_reg_manager):
# 
# Path: morph_net/framework/leaf_op_handler.py
# class LeafOpHandler(op_handler.OpHandler):
#   def is_source_op(self):
#   def is_passthrough(self):
#   def assign_grouping(self, op, op_reg_manager):
#   def create_regularizer(self, _):
# 
# Path: morph_net/framework/output_non_passthrough_op_handler.py
# class OutputNonPassthroughOpHandler(op_handler.OpHandler):
#   def is_source_op(self):
#   def is_passthrough(self):
#   def assign_grouping(self, op, op_reg_manager):
#   def _group_with_output_slices(
#       self, op, output_op_slices, op_slices, op_reg_manager):
#   def _get_output_op_slices(self, output_ops, op_reg_manager):
#   def create_regularizer(self, _):


import collections
from morph_net.framework import concat_op_handler
from morph_net.framework import depthwise_convolution_op_handler
from morph_net.framework import grouping_op_handler
from morph_net.framework import leaf_op_handler
from morph_net.framework import output_non_passthrough_op_handler

def _get_base_op_hander_dicts():
  """Returns the base op_hander_dict for all regularizers."""
  base_dict = collections.defaultdict(
      grouping_op_handler.GroupingOpHandler, {
          'ConcatV2':
              concat_op_handler.ConcatOpHandler(),
          'DepthToSpace':
              output_non_passthrough_op_handler.OutputNonPassthroughOpHandler(),
          'DepthwiseConv2dNative':
              depthwise_convolution_op_handler.DepthwiseConvolutionOpHandler(),
          'ExpandDims':
              output_non_passthrough_op_handler.OutputNonPassthroughOpHandler(),
          'RandomUniform':
              leaf_op_handler.LeafOpHandler(),
          'Reshape':
              leaf_op_handler.LeafOpHandler(),
          'Shape':
              leaf_op_handler.LeafOpHandler(),
          'SpaceToDepth':
              leaf_op_handler.LeafOpHandler(),
          'StridedSlice':
              leaf_op_handler.LeafOpHandler(),
          'TensorArrayGatherV3':
              leaf_op_handler.LeafOpHandler(),
          'Transpose':
              output_non_passthrough_op_handler.OutputNonPassthroughOpHandler(),
      })
  for resize_method in RESIZE_OP_NAMES:
    # Resize* ops, second input might be a tensor which will result in an error.

# Path: morph_net/framework/concat_op_handler.py
# def _get_concat_op_axis(op):
#   def is_source_op(self):
#   def is_passthrough(self):
#   def assign_grouping(self, op, op_reg_manager):
#   def create_regularizer(self, _):
#   def _get_input_output_op_slices(self, input_ops, output_ops, op_reg_manager):
# class ConcatOpHandler(op_handler.OpHandler):
# 
# Path: morph_net/framework/depthwise_convolution_op_handler.py
# class DepthwiseConvolutionOpHandler(grouping_op_handler.GroupingOpHandler):
#   def assign_grouping(self, op, op_reg_manager):
#   def _get_depth_multiplier_input_output_op_slices(
#       self, input_ops, input_size, output_ops, op_reg_manager,
#       depth_multiplier):
# 
# Path: morph_net/framework/grouping_op_handler.py
# class GroupingOpHandler(op_handler.OpHandler):
#   def __init__(self, grouping_indices=None):
#   def is_source_op(self):
#   def is_passthrough(self):
#   def assign_grouping(self, op, op_reg_manager):
#   def create_regularizer(self, _):
#   def _get_input_output_op_slices(self, input_ops, output_ops, op_reg_manager):
#   def _is_broadcast(self, op, op_reg_manager):
# 
# Path: morph_net/framework/leaf_op_handler.py
# class LeafOpHandler(op_handler.OpHandler):
#   def is_source_op(self):
#   def is_passthrough(self):
#   def assign_grouping(self, op, op_reg_manager):
#   def create_regularizer(self, _):
# 
# Path: morph_net/framework/output_non_passthrough_op_handler.py
# class OutputNonPassthroughOpHandler(op_handler.OpHandler):
#   def is_source_op(self):
#   def is_passthrough(self):
#   def assign_grouping(self, op, op_reg_manager):
#   def _group_with_output_slices(
#       self, op, output_op_slices, op_slices, op_reg_manager):
#   def _get_output_op_slices(self, output_ops, op_reg_manager):
#   def create_regularizer(self, _):


# Path: morph_net/framework/op_handlers.py
import collections
from morph_net.framework import concat_op_handler
from morph_net.framework import depthwise_convolution_op_handler
from morph_net.framework import grouping_op_handler
from morph_net.framework import leaf_op_handler
from morph_net.framework import output_non_passthrough_op_handler


def _get_base_op_hander_dicts():
  """Returns the base op_hander_dict for all regularizers."""
  base_dict = collections.defaultdict(
      grouping_op_handler.GroupingOpHandler, {
          'ConcatV2':
              concat_op_handler.ConcatOpHandler(),
          'DepthToSpace':
              output_non_passthrough_op_handler.OutputNonPassthroughOpHandler(),
          'DepthwiseConv2dNative':
              depthwise_convolution_op_handler.DepthwiseConvolutionOpHandler(),
          'ExpandDims':
              output_non_passthrough_op_handler.OutputNonPassthroughOpHandler(),
          'RandomUniform':
              leaf_op_handler.LeafOpHandler(),
          'Reshape':
              leaf_op_handler.LeafOpHandler(),
          'Shape':
              leaf_op_handler.LeafOpHandler(),
          'SpaceToDepth':
              leaf_op_handler.LeafOpHandler(),
          'StridedSlice':
              leaf_op_handler.LeafOpHandler(),
          'TensorArrayGatherV3':
              leaf_op_handler.LeafOpHandler(),
          'Transpose':
              output_non_passthrough_op_handler.OutputNonPassthroughOpHandler(),
      })
  for resize_method in RESIZE_OP_NAMES:
    # Resize* ops, second input might be a tensor which will result in an error.

    base_dict[resize_method] = grouping_op_handler.GroupingOpHandler([0])