# Path: morph_net/framework/logistic_sigmoid_source_op_handler.py
# class LogisticSigmoidSourceOpHandler(grouping_op_handler.GroupingOpHandler):
#   def __init__(self, regularize_on_mask=True,
#                alive_threshold=0.1, mask_as_alive_vector=True):
#   def is_source_op(self):
#   def assign_grouping(self, op, op_reg_manager):
#   def create_regularizer(self, op_slice):
# 
# 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/probabilistic_gating/activation_gating.py
# BN_FN = contrib_layers.batch_norm
# class LogisticSigmoidGating(tf.keras.layers.Layer):
#   def __init__(self,
#                axis=3,
#                temperature=0.001,
#                straight_through=False,
#                log_odds_init=2.5,
#                soft_mask_for_inference=True,
#                keep_first_channel_alive=True,
#                annealing_rate=None,
#                global_step=None,
#                name=None):
#   def build(self, activation_shape):
#   def call(self, activation, is_training):
# def logistic_sigmoid_gating(activation,
#                             axis,
#                             is_training,
#                             temperature=0.001,
#                             straight_through=False,
#                             log_odds_init=2.5,
#                             soft_mask_for_inference=True,
#                             keep_first_channel_alive=True,
#                             annealing_rate=None,
#                             global_step=None,
#                             scope=None):
# def _logistic_sigmoid_sample(logits, tau, straight_through):
#     def _logistic_sigmoid_grad(log_odds, logistic_sample, tau, dy):
#     def _logistic_sigmoid(log_odds, logistic_sample, tau):
# def gated_batch_norm(gating_fn=logistic_sigmoid_gating,
#                      axis=3,
#                      is_training=True,
#                      **kwargs_for_gating_fn):
# def gated_relu_activation(gating_fn=logistic_sigmoid_gating,
#                           axis=3,
#                           is_training=True,
#                           **kwargs_for_gating_fn):
# def add_gating_to_fn(
#     fn, gating_fn=logistic_sigmoid_gating, axis=3,
#     is_training=True, **kwargs_for_gating_fn):
#   def _gated_fn(*args, **kwargs):


from absl.testing import parameterized
from morph_net.framework import logistic_sigmoid_source_op_handler as ls_source_op_handler
from morph_net.framework import op_regularizer_manager as orm
from morph_net.probabilistic_gating import activation_gating
from tensorflow.contrib import framework as contrib_framework
from tensorflow.contrib import layers as contrib_layers
import mock
import scipy as sp
import tensorflow.compat.v1 as tf
        self.activation_gating_op_slice)

    # Verify regularizer is the mask tensor.
    g = tf.get_default_graph()
    mask_tensor = g.get_tensor_by_name(
        'logistic_sigmoid_gating/LogisticSigmoidGating:0')
    self.assertEqual(mask_tensor, regularizer._regularization_vector)

  def testCreateRegularizerOnLogits(self):
    # Call handler to create regularizer.
    handler = ls_source_op_handler.LogisticSigmoidSourceOpHandler(
        regularize_on_mask=False)
    regularizer = handler.create_regularizer(
        self.activation_gating_op_slice)

    # Verify regularizer is the mask tensor.
    g = tf.get_default_graph()
    logits_tensor = g.get_tensor_by_name(
        'logistic_sigmoid_gating/mask_logits/Read/ReadVariableOp:0')

    with self.cached_session() as sess:
      sess.run(tf.global_variables_initializer())
      logits_vec, reg_vec = sess.run(
          [logits_tensor, regularizer._regularization_vector])

      # Verify regularizer is the probability tensor induced by the logits.
      self.assertAllEqual(sp.special.expit(logits_vec), reg_vec)

  def testCreateRegularizer_OnMask_Sliced(self):
    # Call handler to create regularizer.

# Path: morph_net/framework/logistic_sigmoid_source_op_handler.py
# class LogisticSigmoidSourceOpHandler(grouping_op_handler.GroupingOpHandler):
#   def __init__(self, regularize_on_mask=True,
#                alive_threshold=0.1, mask_as_alive_vector=True):
#   def is_source_op(self):
#   def assign_grouping(self, op, op_reg_manager):
#   def create_regularizer(self, op_slice):
# 
# 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/probabilistic_gating/activation_gating.py
# BN_FN = contrib_layers.batch_norm
# class LogisticSigmoidGating(tf.keras.layers.Layer):
#   def __init__(self,
#                axis=3,
#                temperature=0.001,
#                straight_through=False,
#                log_odds_init=2.5,
#                soft_mask_for_inference=True,
#                keep_first_channel_alive=True,
#                annealing_rate=None,
#                global_step=None,
#                name=None):
#   def build(self, activation_shape):
#   def call(self, activation, is_training):
# def logistic_sigmoid_gating(activation,
#                             axis,
#                             is_training,
#                             temperature=0.001,
#                             straight_through=False,
#                             log_odds_init=2.5,
#                             soft_mask_for_inference=True,
#                             keep_first_channel_alive=True,
#                             annealing_rate=None,
#                             global_step=None,
#                             scope=None):
# def _logistic_sigmoid_sample(logits, tau, straight_through):
#     def _logistic_sigmoid_grad(log_odds, logistic_sample, tau, dy):
#     def _logistic_sigmoid(log_odds, logistic_sample, tau):
# def gated_batch_norm(gating_fn=logistic_sigmoid_gating,
#                      axis=3,
#                      is_training=True,
#                      **kwargs_for_gating_fn):
# def gated_relu_activation(gating_fn=logistic_sigmoid_gating,
#                           axis=3,
#                           is_training=True,
#                           **kwargs_for_gating_fn):
# def add_gating_to_fn(
#     fn, gating_fn=logistic_sigmoid_gating, axis=3,
#     is_training=True, **kwargs_for_gating_fn):
#   def _gated_fn(*args, **kwargs):


# Path: morph_net/framework/logistic_sigmoid_source_op_handler_test.py
from absl.testing import parameterized
from morph_net.framework import logistic_sigmoid_source_op_handler as ls_source_op_handler
from morph_net.framework import op_regularizer_manager as orm
from morph_net.probabilistic_gating import activation_gating
from tensorflow.contrib import framework as contrib_framework
from tensorflow.contrib import layers as contrib_layers
import mock
import scipy as sp
import tensorflow.compat.v1 as tf

        self.activation_gating_op_slice)

    # Verify regularizer is the mask tensor.
    g = tf.get_default_graph()
    mask_tensor = g.get_tensor_by_name(
        'logistic_sigmoid_gating/LogisticSigmoidGating:0')
    self.assertEqual(mask_tensor, regularizer._regularization_vector)

  def testCreateRegularizerOnLogits(self):
    # Call handler to create regularizer.
    handler = ls_source_op_handler.LogisticSigmoidSourceOpHandler(
        regularize_on_mask=False)
    regularizer = handler.create_regularizer(
        self.activation_gating_op_slice)

    # Verify regularizer is the mask tensor.
    g = tf.get_default_graph()
    logits_tensor = g.get_tensor_by_name(
        'logistic_sigmoid_gating/mask_logits/Read/ReadVariableOp:0')

    with self.cached_session() as sess:
      sess.run(tf.global_variables_initializer())
      logits_vec, reg_vec = sess.run(
          [logits_tensor, regularizer._regularization_vector])

      # Verify regularizer is the probability tensor induced by the logits.
      self.assertAllEqual(sp.special.expit(logits_vec), reg_vec)

  def testCreateRegularizer_OnMask_Sliced(self):
    # Call handler to create regularizer.

    handler = ls_source_op_handler.LogisticSigmoidSourceOpHandler(