# 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/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 mock
import tensorflow.compat.v1 as tf
from morph_net.framework import op_regularizer_manager as orm
from morph_net.framework import output_non_passthrough_op_handler
from tensorflow.contrib import framework
from tensorflow.contrib import layers


class OutputNonPassthroughOpHandlerTest(tf.test.TestCase):

  def _batch_norm_scope(self):
    params = {
        'trainable': True,
        'normalizer_fn': layers.batch_norm,
        'normalizer_params': {
            'scale': True,
        },
    }

    with framework.arg_scope([layers.conv2d], **params) as sc:
      return sc

  def setUp(self):
    super(OutputNonPassthroughOpHandlerTest, self).setUp()

    # This tests 2 Conv2D ops with batch norm at the top.
    with framework.arg_scope(self._batch_norm_scope()):
      inputs = tf.zeros([2, 4, 4, 3])
      c1 = layers.conv2d(inputs, num_outputs=5,
                         kernel_size=3, scope='conv1', normalizer_fn=None)
      layers.conv2d(c1, num_outputs=6, kernel_size=3, scope='conv2')

    g = tf.get_default_graph()

    # Declare OpSlice and OpGroup for ops of interest.
    self.conv1_op = g.get_operation_by_name('conv1/Conv2D')

# 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/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/output_non_passthrough_op_handler_test.py
import mock
import tensorflow.compat.v1 as tf
from morph_net.framework import op_regularizer_manager as orm
from morph_net.framework import output_non_passthrough_op_handler
from tensorflow.contrib import framework
from tensorflow.contrib import layers



class OutputNonPassthroughOpHandlerTest(tf.test.TestCase):

  def _batch_norm_scope(self):
    params = {
        'trainable': True,
        'normalizer_fn': layers.batch_norm,
        'normalizer_params': {
            'scale': True,
        },
    }

    with framework.arg_scope([layers.conv2d], **params) as sc:
      return sc

  def setUp(self):
    super(OutputNonPassthroughOpHandlerTest, self).setUp()

    # This tests 2 Conv2D ops with batch norm at the top.
    with framework.arg_scope(self._batch_norm_scope()):
      inputs = tf.zeros([2, 4, 4, 3])
      c1 = layers.conv2d(inputs, num_outputs=5,
                         kernel_size=3, scope='conv1', normalizer_fn=None)
      layers.conv2d(c1, num_outputs=6, kernel_size=3, scope='conv2')

    g = tf.get_default_graph()

    # Declare OpSlice and OpGroup for ops of interest.
    self.conv1_op = g.get_operation_by_name('conv1/Conv2D')

    self.conv1_op_slice = orm.OpSlice(self.conv1_op, orm.Slice(0, 5))