# 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/tpu_util.py
# def run_on_cpu():
# def is_on_cpu():
# def get_variable_name(read_variable_op):
# def maybe_convert_to_variable(tensor):
# def write_to_variable(tensor, fail_if_exists=True):
# def read_from_variable(tensor):
# def is_intermediate_var(v):
# 
# Path: morph_net/op_regularizers/gamma_l1_regularizer.py
# class GammaL1Regularizer(generic_regularizers.OpRegularizer):
#   def __init__(self, gamma, gamma_threshold):
#   def regularization_vector(self):
#   def alive_vector(self):


from morph_net.framework import grouping_op_handler
from morph_net.framework import tpu_util
from morph_net.op_regularizers import gamma_l1_regularizer
    op_slices = op_reg_manager.get_op_slices(op)
    for op_slice in op_slices:
      op_group = op_reg_manager.get_op_group(op_slice)
      if op_group is None:
        op_reg_manager.create_op_group_for_op_slice(op_slice)

    super(BatchNormSourceOpHandler, self).assign_grouping(op, op_reg_manager)

  def create_regularizer(self, op_slice):
    """Create a regularizer for this batch norm OpSlice.

    Args:
      op_slice: op_regularizer_manager.OpSlice that is a batch norm OpSlice.

    Returns:
      OpRegularizer for this batch norm op.
    """
    start_index = op_slice.slice.start_index
    size = op_slice.slice.size
    gamma = op_slice.op.inputs[1]  # Input 1 is gamma.

    # If OpSlice size matches tensor size, use the entire tensor.  Otherwise,
    # slice the tensor accordingly.
    if start_index == 0 and size == gamma.shape.as_list()[-1]:
      return gamma_l1_regularizer.GammaL1Regularizer(
          gamma, self._gamma_threshold)
    else:
      # Note: this conversion is also attempted inside GammaL1Regularizer
      # because it may be invoked from another call site.
      gamma = tpu_util.maybe_convert_to_variable(gamma)

# 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/tpu_util.py
# def run_on_cpu():
# def is_on_cpu():
# def get_variable_name(read_variable_op):
# def maybe_convert_to_variable(tensor):
# def write_to_variable(tensor, fail_if_exists=True):
# def read_from_variable(tensor):
# def is_intermediate_var(v):
# 
# Path: morph_net/op_regularizers/gamma_l1_regularizer.py
# class GammaL1Regularizer(generic_regularizers.OpRegularizer):
#   def __init__(self, gamma, gamma_threshold):
#   def regularization_vector(self):
#   def alive_vector(self):


# Path: morph_net/framework/batch_norm_source_op_handler.py
from morph_net.framework import grouping_op_handler
from morph_net.framework import tpu_util
from morph_net.op_regularizers import gamma_l1_regularizer

    op_slices = op_reg_manager.get_op_slices(op)
    for op_slice in op_slices:
      op_group = op_reg_manager.get_op_group(op_slice)
      if op_group is None:
        op_reg_manager.create_op_group_for_op_slice(op_slice)

    super(BatchNormSourceOpHandler, self).assign_grouping(op, op_reg_manager)

  def create_regularizer(self, op_slice):
    """Create a regularizer for this batch norm OpSlice.

    Args:
      op_slice: op_regularizer_manager.OpSlice that is a batch norm OpSlice.

    Returns:
      OpRegularizer for this batch norm op.
    """
    start_index = op_slice.slice.start_index
    size = op_slice.slice.size
    gamma = op_slice.op.inputs[1]  # Input 1 is gamma.

    # If OpSlice size matches tensor size, use the entire tensor.  Otherwise,
    # slice the tensor accordingly.
    if start_index == 0 and size == gamma.shape.as_list()[-1]:
      return gamma_l1_regularizer.GammaL1Regularizer(
          gamma, self._gamma_threshold)
    else:
      # Note: this conversion is also attempted inside GammaL1Regularizer
      # because it may be invoked from another call site.
      gamma = tpu_util.maybe_convert_to_variable(gamma)

      return gamma_l1_regularizer.GammaL1Regularizer(