# 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/op_regularizer_manager.py
# ITERATION_LIMIT = 1000000
# class OpSlice(collections.namedtuple('OpSlice', ['op', 'slice'])):
# class Slice(collections.namedtuple('Slice', ['start_index', 'size'])):
# class OpRegularizerManager(object):
# class OpGroup(object):
#   def __str__(self):
#   def __str__(self):
#   def __init__(
#       self,
#       output_boundary: List[tf.Operation],
#       op_handler_dict=None,
#       create_grouping_regularizer=grouping_regularizers.MaxGroupingRegularizer,
#       force_group=None,
#       regularizer_blacklist=None,
#       input_boundary: List[tf.Operation] = None,
#       iteration_limit=ITERATION_LIMIT):
#   def ops(self):
#   def get_regularizer(self, op):
#   def create_op_group_for_op_slice(self, op_slice, is_source=True):
#   def group_op_slices(self, op_slices, omit_source_op_slices=None):
#   def slice_op(self, op, sizes):
#   def process_ops(self, ops):
#   def process_ops_last(self, ops):
#   def is_source_op(self, op):
#   def is_passthrough(self, op):
#   def get_op_slices(self, op):
#   def get_op_group(self, op_slice):
#   def _slice_op_slice(self, op_slice, sizes, size_index, size_count,
#                       new_op_slice_group):
#   def _slice_op_with_sizes(self, op, sizes, is_source, is_resliced):
#   def _get_source_slices(self, sizes, op_slices):
#   def _dfs_for_source_ops(self, output_boundary, input_boundary=None):
#   def _force_group_ops(self, force_group):
#   def __init__(self, op_slice=None, op_groups=None, omit_source_op_slices=None):
#   def op_slices(self):
#   def source_op_slices(self):


import mock
import tensorflow.compat.v1 as tf
from morph_net.framework import concat_op_handler
from morph_net.framework import op_regularizer_manager as orm
from tensorflow.contrib import framework
from tensorflow.contrib import layers
    self.mock_op_reg_manager.group_op_slices.assert_has_calls(
        [mock.call([concat_op_slice_0_5, self.relu1_op_slice,
                    batch_norm_op_slice_0_5]),
         mock.call([concat_op_slice_5_9, relu2_op_slice_0_4,
                    batch_norm_op_slice_5_9]),
         mock.call([concat_op_slice_9_11, relu2_op_slice_4_6,
                    batch_norm_op_slice_9_11]),
         mock.call([concat_op_slice_11_13, relu3_op_slice_0_2,
                    batch_norm_op_slice_11_13]),
         mock.call([concat_op_slice_13_18, relu3_op_slice_2_7,
                    batch_norm_op_slice_13_18])])

  def testAssignGrouping_AllNeighborsGrouped_InputSlicesNotAligned(self):
    # In this test, the op c2 has size 6 but is split into 2 slices of size 3.
    # The concat op (and its output, the batch norm) both have size 18.  This
    # test verifies that the concat and batch norm are sliced according to the
    # sizes of c1, c2, and c3, and takes into account that c2 is also composed
    # of multiple slices.
    concat_op_slice_0_5 = orm.OpSlice(self.concat_op, orm.Slice(0, 5))
    concat_op_slice_5_8 = orm.OpSlice(self.concat_op, orm.Slice(5, 3))
    concat_op_slice_8_11 = orm.OpSlice(self.concat_op, orm.Slice(8, 3))
    concat_op_slice_11_18 = orm.OpSlice(self.concat_op, orm.Slice(11, 7))

    relu2_op_slice_0_3 = orm.OpSlice(self.relu2_op, orm.Slice(0, 3))
    relu2_op_slice_3_6 = orm.OpSlice(self.relu2_op, orm.Slice(3, 3))
    relu2_op_group1 = orm.OpGroup(
        relu2_op_slice_0_3, omit_source_op_slices=[relu2_op_slice_0_3])
    relu2_op_group2 = orm.OpGroup(
        relu2_op_slice_3_6, omit_source_op_slices=[relu2_op_slice_3_6])


# 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/op_regularizer_manager.py
# ITERATION_LIMIT = 1000000
# class OpSlice(collections.namedtuple('OpSlice', ['op', 'slice'])):
# class Slice(collections.namedtuple('Slice', ['start_index', 'size'])):
# class OpRegularizerManager(object):
# class OpGroup(object):
#   def __str__(self):
#   def __str__(self):
#   def __init__(
#       self,
#       output_boundary: List[tf.Operation],
#       op_handler_dict=None,
#       create_grouping_regularizer=grouping_regularizers.MaxGroupingRegularizer,
#       force_group=None,
#       regularizer_blacklist=None,
#       input_boundary: List[tf.Operation] = None,
#       iteration_limit=ITERATION_LIMIT):
#   def ops(self):
#   def get_regularizer(self, op):
#   def create_op_group_for_op_slice(self, op_slice, is_source=True):
#   def group_op_slices(self, op_slices, omit_source_op_slices=None):
#   def slice_op(self, op, sizes):
#   def process_ops(self, ops):
#   def process_ops_last(self, ops):
#   def is_source_op(self, op):
#   def is_passthrough(self, op):
#   def get_op_slices(self, op):
#   def get_op_group(self, op_slice):
#   def _slice_op_slice(self, op_slice, sizes, size_index, size_count,
#                       new_op_slice_group):
#   def _slice_op_with_sizes(self, op, sizes, is_source, is_resliced):
#   def _get_source_slices(self, sizes, op_slices):
#   def _dfs_for_source_ops(self, output_boundary, input_boundary=None):
#   def _force_group_ops(self, force_group):
#   def __init__(self, op_slice=None, op_groups=None, omit_source_op_slices=None):
#   def op_slices(self):
#   def source_op_slices(self):


# Path: morph_net/framework/concat_op_handler_test.py
import mock
import tensorflow.compat.v1 as tf
from morph_net.framework import concat_op_handler
from morph_net.framework import op_regularizer_manager as orm
from tensorflow.contrib import framework
from tensorflow.contrib import layers

    self.mock_op_reg_manager.group_op_slices.assert_has_calls(
        [mock.call([concat_op_slice_0_5, self.relu1_op_slice,
                    batch_norm_op_slice_0_5]),
         mock.call([concat_op_slice_5_9, relu2_op_slice_0_4,
                    batch_norm_op_slice_5_9]),
         mock.call([concat_op_slice_9_11, relu2_op_slice_4_6,
                    batch_norm_op_slice_9_11]),
         mock.call([concat_op_slice_11_13, relu3_op_slice_0_2,
                    batch_norm_op_slice_11_13]),
         mock.call([concat_op_slice_13_18, relu3_op_slice_2_7,
                    batch_norm_op_slice_13_18])])

  def testAssignGrouping_AllNeighborsGrouped_InputSlicesNotAligned(self):
    # In this test, the op c2 has size 6 but is split into 2 slices of size 3.
    # The concat op (and its output, the batch norm) both have size 18.  This
    # test verifies that the concat and batch norm are sliced according to the
    # sizes of c1, c2, and c3, and takes into account that c2 is also composed
    # of multiple slices.
    concat_op_slice_0_5 = orm.OpSlice(self.concat_op, orm.Slice(0, 5))
    concat_op_slice_5_8 = orm.OpSlice(self.concat_op, orm.Slice(5, 3))
    concat_op_slice_8_11 = orm.OpSlice(self.concat_op, orm.Slice(8, 3))
    concat_op_slice_11_18 = orm.OpSlice(self.concat_op, orm.Slice(11, 7))

    relu2_op_slice_0_3 = orm.OpSlice(self.relu2_op, orm.Slice(0, 3))
    relu2_op_slice_3_6 = orm.OpSlice(self.relu2_op, orm.Slice(3, 3))
    relu2_op_group1 = orm.OpGroup(
        relu2_op_slice_0_3, omit_source_op_slices=[relu2_op_slice_0_3])
    relu2_op_group2 = orm.OpGroup(
        relu2_op_slice_3_6, omit_source_op_slices=[relu2_op_slice_3_6])


    batch_norm_op_slice = orm.OpSlice(self.batch_norm_op, orm.Slice(0, 18))