# 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/prob_gating_regularizer.py
# class ProbGatingRegularizer(generic_regularizers.OpRegularizer):
#   def __init__(self, logits, mask, regularize_on_mask=True,
#                alive_threshold=0.1, mask_as_alive_vector=True):
#   def regularization_vector(self):
#   def alive_vector(self):
#   def is_probabilistic(self):


from morph_net.framework import grouping_op_handler
from morph_net.framework import tpu_util
from morph_net.op_regularizers import prob_gating_regularizer
        op_reg_manager.create_op_group_for_op_slice(op_slice)

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

  def create_regularizer(self, op_slice):
    """Create a regularizer for this logistic-sigmoid gating 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
    logits = op_slice.op.inputs[0]  # Input 0 are the logits.
    mask = op_slice.op.outputs[0]

    # If OpSlice size matches tensor size, use the entire tensor.  Otherwise,
    # slice the tensor accordingly.
    mask = tpu_util.read_from_variable(mask)

    if start_index == 0 and size == logits.shape.as_list()[-1]:
      return prob_gating_regularizer.ProbGatingRegularizer(
          logits, mask,
          regularize_on_mask=self._regularize_on_mask,
          alive_threshold=self._alive_threshold,
          mask_as_alive_vector=self._mask_as_alive_vector)
    else:

# 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/prob_gating_regularizer.py
# class ProbGatingRegularizer(generic_regularizers.OpRegularizer):
#   def __init__(self, logits, mask, regularize_on_mask=True,
#                alive_threshold=0.1, mask_as_alive_vector=True):
#   def regularization_vector(self):
#   def alive_vector(self):
#   def is_probabilistic(self):


# Path: morph_net/framework/logistic_sigmoid_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 prob_gating_regularizer

        op_reg_manager.create_op_group_for_op_slice(op_slice)

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

  def create_regularizer(self, op_slice):
    """Create a regularizer for this logistic-sigmoid gating 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
    logits = op_slice.op.inputs[0]  # Input 0 are the logits.
    mask = op_slice.op.outputs[0]

    # If OpSlice size matches tensor size, use the entire tensor.  Otherwise,
    # slice the tensor accordingly.
    mask = tpu_util.read_from_variable(mask)

    if start_index == 0 and size == logits.shape.as_list()[-1]:
      return prob_gating_regularizer.ProbGatingRegularizer(
          logits, mask,
          regularize_on_mask=self._regularize_on_mask,
          alive_threshold=self._alive_threshold,
          mask_as_alive_vector=self._mask_as_alive_vector)
    else:

      logits = tpu_util.maybe_convert_to_variable(logits)