class BatchNormSourceOpHandler(grouping_op_handler.GroupingOpHandler):
  def __init__(self, gamma_threshold):
  def is_source_op(self):
  def assign_grouping(self, op, op_reg_manager):
  def create_regularizer(self, op_slice): 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): class ConvSourceOpHandler(
    group_lasso_base_op_handler.GroupLassoBaseSourceOpHandler):
  def _reduce_dims(self, op): 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): class NetworkRegularizer(object):  # pytype: disable=ignored-metaclass
class OpRegularizer(object):  # pytype: disable=ignored-metaclass
  def get_regularization_term(self, ops=None):
  def get_cost(self, ops=None):
  def op_regularizer_manager(self):
  def name(self):
  def cost_name(self):
  def regularization_vector(self):
  def alive_vector(self):
def dimensions_are_compatible(op_regularizer): 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): 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): 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, _): ALIVE_STUB = {
    'conv1': [False, True, True, False, True, False, True],
    'conv2': [True, False, True, False, False],
    'conv3': [False, False, True, True],
    'conv4': [False, True, False, True, True, False, True,
              False, False, True, False, False],
    'conv5': [False, True, False],
}
REG_STUB = {
    'conv1': [0.1, 0.3, 0.6, 0.2, 0.4, 0.0, 0.8],
    'conv2': [0.15, 0.25, 0.05, 0.55, 0.45],
    'conv3': [0.07, 0.27, 0.17, 0.37],
    'conv4': [
        0.07, 0.27, 0.17, 0.37, 0.28, 0.32, 0.12, 0.22, 0.19, 0.11, 0.02, 0.06
    ],
    'conv5': [0.24, 0.34, 0.29],
}
MOCK_REG_DICT = {
    'Conv2D':
        op_regularizer_stub.OpRegularizerStubFactory(ALIVE_STUB, REG_STUB)
        .create_conv2d_regularizer
}
def image_stub():
def build_model(image=None):
def expected_regularization():
def expected_alive(): ALIVE_STUB = {
    'conv1': [False, True, True],
    'conv2': [True, False, True],
    'conv3': [False, False, True, True],
    'conv4': [False, True, False, True],
}
REG_STUB = {
    'conv1': [0.1, 0.3, 0.6],
    'conv2': [0.15, 0.25, 0.05],
    'conv3': [0.07, 0.17, 0.57, 0.37],
    'conv4': [0.06, 100.87, 0.17, 1.57],
}
MOCK_REG_DICT = {
    'Conv2D':
        op_regularizer_stub.OpRegularizerStubFactory(ALIVE_STUB, REG_STUB)
        .create_conv2d_regularizer
}
def image_stub():
def build_model(image=None):
def expected_regularization():
def expected_alive():
import collections
import numpy as np
import tensorflow.compat.v1 as tf
from absl.testing import parameterized
from morph_net.framework import batch_norm_source_op_handler
from morph_net.framework import concat_op_handler
from morph_net.framework import conv_source_op_handler
from morph_net.framework import depthwise_convolution_op_handler
from morph_net.framework import generic_regularizers
from morph_net.framework import grouping_op_handler
from morph_net.framework import op_regularizer_manager as orm
from morph_net.framework import output_non_passthrough_op_handler
from morph_net.testing import add_concat_model_stub
from morph_net.testing import grouping_concat_model_stub
from tensorflow.contrib import framework
from tensorflow.contrib import layers

    # All ops are in the same group
    group = list(op_reg_manager._op_group_dict.values())[0]
    source_op_names = [s.op.name for s in group.source_op_slices]
    self.assertSetEqual(set(['bn0/FusedBatchNormV3', 'bn1/FusedBatchNormV3']),
                        set(source_op_names))

  def testPrintOpSlices(self):
    inputs = tf.zeros([2, 4, 4, 3])
    identity1 = tf.identity(inputs)
    identity2 = tf.identity(inputs)

    manager = orm.OpRegularizerManager(
        [identity1.op, identity2.op],
        op_handler_dict=self._default_op_handler_dict)
    op_slices1 = manager.get_op_slices(identity1.op)
    op_slices2 = manager.get_op_slices(identity2.op)
    all_slices = op_slices1 + op_slices2

    self.assertEqual('[Identity (0, 3), Identity_1 (0, 3)]',
                     str(all_slices))


class IndexOpRegularizer(generic_regularizers.OpRegularizer):
  """A test OpRegularizer with a self-incrementing index.

  This class creates a regularizer where the regularization vector contains
  self-incrementing values (e.g. [0, 1, 2, ...]).  The index continues to
  increment as regularizers are created.  This is convenient for testing in
  order to track individual elements of the regularization vector (e.g. gather).

  For example, creating 2 regularizers of size 3 results in
  r1 = [0, 1, 2] and r2 = [3, 4, 5].
  """

  index = 0

  def __init__(self, op_slice, op_reg_manager):
    size = op_slice.slice.size
    self._alive_vector = tf.cast(tf.ones(size), tf.bool)
    self._regularization_vector = tf.constant(
        list(range(IndexOpRegularizer.index, IndexOpRegularizer.index + size)),
        tf.float32)
    IndexOpRegularizer.index += size

  @classmethod
  def reset_index(cls):
    IndexOpRegularizer.index = 0

  @property
  def regularization_vector(self):
    return self._regularization_vector

  @property
  def alive_vector(self):
    return self._alive_vector


class SumGroupingRegularizer(generic_regularizers.OpRegularizer):
  """A regularizer that groups others by summing their regularization values."""

  def __init__(self, regularizers_to_group):
    """Creates an instance.

    Args:
      regularizers_to_group: A list of generic_regularizers.OpRegularizer
        objects.Their regularization_vector (alive_vector) are expected to be of
        the same length.

    Raises:
      ValueError: regularizers_to_group is not of length at least 2.
    """
    if len(regularizers_to_group) < 2:
      raise ValueError('Groups must be of at least size 2.')
    self._regularization_vector = tf.add_n(
        [r.regularization_vector for r in regularizers_to_group])
    self._alive_vector = tf.cast(
        tf.ones(self._regularization_vector.get_shape()[-1]), tf.bool)

  @property
  def regularization_vector(self):
    return self._regularization_vector

  @property
  def alive_vector(self):
    return self._alive_vector


class IndexBatchNormSourceOpHandler(
    batch_norm_source_op_handler.BatchNormSourceOpHandler):
  """An OpHandler that creates OpRegularizer using IndexOpRegularizer.

  A wrapper around BatchNormSourceOpHandler that overrides the
  create_regularizer method to use IndexOpRegularizer for testing.
  """

  def __init__(self):
    super(IndexBatchNormSourceOpHandler, self).__init__(0.0)

  def create_regularizer(self, op_slice):
    return IndexOpRegularizer(op_slice, None)


class StubBatchNormSourceOpHandler(
    batch_norm_source_op_handler.BatchNormSourceOpHandler):
  """An OpHandler that creates OpRegularizer using stub values.

  A wrapper around BatchNormSourceOpHandler that overrides the
  create_regularizer method to use stub values for testing.
  """

  def __init__(self, model_stub):
    super(StubBatchNormSourceOpHandler, self).__init__(0.0)
    self._model_stub = model_stub

  def create_regularizer(self, op_slice):
    return _stub_create_regularizer(op_slice, self._model_stub)


class IndexConvSourceOpHandler(

    conv_source_op_handler.ConvSourceOpHandler):