# 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.probabilistic_gating import activation_gating
from tensorflow.contrib import layers as contrib_layers
import numpy as np
import tensorflow.compat.v1 as tf
    mask_shape = op.outputs[0].shape.as_list()
    self.assertLen(mask_shape, 1)
    self.assertEqual(mask_shape[0], activation.shape.as_list()[axis])

  @parameterized.named_parameters(('_1', 5.5, 0.996, True),
                                  ('_2', 0.5, 0.6, True),
                                  ('_3', -5.5, 0.004, True),
                                  ('_4', 5.5, 0.996, False),
                                  ('_5', 0.5, 0.6, False),
                                  ('_6', -5.5, 0.004, False),
                                  ('_7', -1000.0, 0.0, False),
                                  ('_8', 1000.0, 1.0, False))
  def test_sampled_mask(self, logits, expected_percent_on, straight_through):
    logits = tf.constant(logits, shape=[1000])
    mask = ag._logistic_sigmoid_sample(
        logits, 0.01, straight_through=straight_through)

    with self.session() as sess:
      sess.run(tf.global_variables_initializer())
      mask = mask.eval()

    avg_mask = np.average(mask)
    self.assertAlmostEqual(avg_mask, expected_percent_on, delta=0.07)

  @parameterized.named_parameters(('_st', True),
                                  ('_no_st', False))
  def test_add_gating_to_fn(self, straight_through):
    activation = tf.reshape(tf.eye(16, 16), shape=[1, 16, 16, 1])
    activation = tf.concat([activation, activation, activation], axis=3)


# 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/probabilistic_gating/activation_gating_test.py
from absl.testing import parameterized
from morph_net.probabilistic_gating import activation_gating
from tensorflow.contrib import layers as contrib_layers
import numpy as np
import tensorflow.compat.v1 as tf

    mask_shape = op.outputs[0].shape.as_list()
    self.assertLen(mask_shape, 1)
    self.assertEqual(mask_shape[0], activation.shape.as_list()[axis])

  @parameterized.named_parameters(('_1', 5.5, 0.996, True),
                                  ('_2', 0.5, 0.6, True),
                                  ('_3', -5.5, 0.004, True),
                                  ('_4', 5.5, 0.996, False),
                                  ('_5', 0.5, 0.6, False),
                                  ('_6', -5.5, 0.004, False),
                                  ('_7', -1000.0, 0.0, False),
                                  ('_8', 1000.0, 1.0, False))
  def test_sampled_mask(self, logits, expected_percent_on, straight_through):
    logits = tf.constant(logits, shape=[1000])
    mask = ag._logistic_sigmoid_sample(
        logits, 0.01, straight_through=straight_through)

    with self.session() as sess:
      sess.run(tf.global_variables_initializer())
      mask = mask.eval()

    avg_mask = np.average(mask)
    self.assertAlmostEqual(avg_mask, expected_percent_on, delta=0.07)

  @parameterized.named_parameters(('_st', True),
                                  ('_no_st', False))
  def test_add_gating_to_fn(self, straight_through):
    activation = tf.reshape(tf.eye(16, 16), shape=[1, 16, 16, 1])
    activation = tf.concat([activation, activation, activation], axis=3)


    gated_bn = activation_gating.gated_batch_norm(