refactored and linted

Author: gisilvs

tensorflow_probability/python/experimental/bijectors/highway_flow_test.py

@@ -1,3 +1,17 @@
+# Copyright 2021 The TensorFlow Probability Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
 """Tests for HighwayFlow."""
 import tensorflow.compat.v2 as tf
 
@@ -8,32 +22,6 @@
 tfb = tfp.bijectors
 tfd = tfp.distributions
 
-#FIXME: test_util.test_seed throws an error
-seed = 1  # test_util.test_seed(sampler_type='stateless')
-
-
-def _dx(x, activation):
-  if activation == 'sigmoid':
-    return tf.math.sigmoid(x) * (1 - tf.math.sigmoid(x))
-  elif activation == 'softplus':
-    return tf.math.sigmoid(x)
-  elif activation == 'tanh':
-    return 1. - tf.math.tanh(x) ** 2
-
-
-def _activation_log_det_jacobian(x, residual_fraction, activation, width,
-                                 gate_first_n):
-  if activation == 'none':
-    return tf.zeros(x.shape[0])
-  else:
-    return tf.reduce_sum(tf.math.log(
-      tf.concat([(residual_fraction) * tf.ones(
-        gate_first_n), tf.zeros(width - gate_first_n)],
-                axis=0) + tf.concat([(1. - residual_fraction) * tf.ones(
-        gate_first_n), tf.ones(width - gate_first_n)],
-                                    axis=0) * _dx(x, activation)),
-      -1)
-
 
 @test_util.test_all_tf_execution_regimes
 class HighwayFlowTests(test_util.TestCase):
@@ -45,13 +33,16 @@ def testBijector(self):
         # Test generic case with scalar input
         x = tf.ones((width,)) * samplers.uniform((width,), minval=-1.,
                                                  maxval=1.,
-                                                 seed=seed)
+                                                 seed=test_util.test_seed(
+                                                   sampler_type='stateless'))
       elif dim == 1:
         # Test with 2D tensor + batch
         x = tf.ones((5, width,
                      width)) * samplers.uniform((5, width, width),
                                                 minval=-1.,
-                                                maxval=1., seed=seed)
+                                                maxval=1.,
+                                                seed=test_util.test_seed(
+                                                  sampler_type='stateless'))
 
       bijector = tfp.experimental.bijectors.build_highway_flow_layer(
         width, activation_fn=True)
@@ -65,12 +56,34 @@ def testBijector(self):
         -bijector.inverse_log_det_jacobian(
           tf.identity(bijector.forward(x)), event_ndims=dim + 1))
 
+  def testBijectorWithoutActivation(self):
+    width = 4
+    x = tf.ones(2, width,
+                width) * samplers.uniform((2, width, width),
+                                          minval=-1.,
+                                          maxval=1.,
+                                          seed=test_util.test_seed(
+                                            sampler_type='stateless'))
+
+    bijector = tfp.experimental.bijectors.build_highway_flow_layer(
+      width, activation_fn=False)
+    self.evaluate(
+      [v.initializer for v in bijector.trainable_variables])
+    self.assertStartsWith(bijector.name, 'highway_flow')
+    self.assertAllClose(x, bijector.inverse(
+      tf.identity(bijector.forward(x))))
+    self.assertAllClose(
+      bijector.forward_log_det_jacobian(x, event_ndims=2),
+      -bijector.inverse_log_det_jacobian(
+        tf.identity(bijector.forward(x)), event_ndims=2))
+
   def testGating(self):
     width = 4
     x = tf.ones((2, width,
                  width)) * samplers.uniform((2, width, width),
                                             minval=-1.,
-                                            maxval=1., seed=seed)
+                                            maxval=1., seed=test_util.test_seed(
+        sampler_type='stateless'))
 
     # Test with gating half of the inputs
     bijector = tfp.experimental.bijectors.build_highway_flow_layer(
@@ -98,66 +111,6 @@ def testGating(self):
       -bijector.inverse_log_det_jacobian(
         tf.identity(bijector.forward(x)), event_ndims=2))
 
-  def testJacobianWithActivation(self):
-    activations = ['softplus']
-    batch_size = 3
-    width = 4
-    dtype = tf.float32
-    gate_first_n = 2
-    residual_fraction = tf.constant(0.5)
-    for activation in activations:
-
-      if activation == 'sigmoid':
-        activation_fn = tf.nn.sigmoid
-      elif activation == 'softplus':
-        activation_fn = tf.nn.softplus
-      elif activation == 'tanh':
-        activation_fn = tf.nn.tanh
-      elif activation == 'none':
-        activation_fn = None
-
-      bijector = tfp.experimental.bijectors.HighwayFlow(
-        residual_fraction=residual_fraction,
-        activation_fn=activation_fn,
-        bias=tf.zeros(width),
-        upper_diagonal_weights_matrix=tf.eye(width),
-        lower_diagonal_weights_matrix=tf.eye(width),
-        gate_first_n=gate_first_n,
-      )
-
-      self.evaluate(
-        [v.initializer for v in bijector.trainable_variables])
-      x = tf.ones((batch_size,
-                   width)) * samplers.uniform((batch_size, width), -10.,
-                                              10., seed=seed)
-      if activation == 'none':
-        y = x
-      else:
-        y = tf.concat([(residual_fraction) * tf.ones(gate_first_n),
-                       tf.zeros(width - gate_first_n)],
-                      axis=0) * x + tf.concat(
-          [(1. - residual_fraction) * tf.ones(
-            gate_first_n), tf.ones(width - gate_first_n)],
-          axis=0) * activation_fn(x)
-      expected_forward_log_det_jacobian = \
-        _activation_log_det_jacobian(x,
-                                     residual_fraction,
-                                     activation,
-                                     width,
-                                     gate_first_n)
-      expected_inverse_log_det_jacobian = \
-        -expected_forward_log_det_jacobian
-      self.assertAllClose(y, bijector.forward(x))
-      self.assertAllClose(x, bijector.inverse(y))
-      self.assertAllClose(
-        expected_inverse_log_det_jacobian,
-        bijector.inverse_log_det_jacobian(y, event_ndims=1),
-      )
-      self.assertAllClose(
-        expected_forward_log_det_jacobian,
-        bijector.forward_log_det_jacobian(x, event_ndims=1),
-      )
-
   def testResidualFractionGradientsWithCenteredDifference(self):
     width = 4
     batch_size = 3
@@ -180,10 +133,12 @@ def testResidualFractionGradientsWithCenteredDifference(self):
 
     h = 1e-3
 
+    # pylint: disable=protected-access
     bijector._residual_fraction = residual_fraction + h
     y1 = tf.reduce_mean(target.log_prob(bijector.forward(x)))
     bijector._residual_fraction = residual_fraction - h
     y2 = tf.reduce_mean(target.log_prob(bijector.forward(x)))
+    # pylint: enable=protected-access
 
     manual_grad = (y1 - y2) / (2 * h)