updated imports

Author: gisilvs

tensorflow_probability/python/experimental/bijectors/BUILD

@@ -109,12 +109,22 @@ multi_substrate_py_library(
     srcs = ["highway_flow.py"],
     srcs_version = "PY3",
     deps = [
-        ":scalar_function_with_inferred_inverse",
-        # numpy dep,
         # tensorflow dep,
-        "//tensorflow_probability/python/bijectors",
-        "//tensorflow_probability/python/util",
+        "//tensorflow_probability/python/bijectors:bijector",
+        "//tensorflow_probability/python/bijectors:chain",
+        "//tensorflow_probability/python/bijectors:fill_scale_tril",
+        "//tensorflow_probability/python/bijectors:fill_triangular",
+        "//tensorflow_probability/python/bijectors:pad",
+        "//tensorflow_probability/python/bijectors:shift",
+        "//tensorflow_probability/python/bijectors:sigmoid",
+        "//tensorflow_probability/python/bijectors:softplus",
+        "//tensorflow_probability/python/bijectors:transform_diagonal",
+        "//tensorflow_probability/python/internal:cache_util",
+        "//tensorflow_probability/python/internal:dtype_util",
+        "//tensorflow_probability/python/internal:prefer_static",
         "//tensorflow_probability/python/internal:samplers",
+        "//tensorflow_probability/python/internal:tensor_util",
+        "//tensorflow_probability/python/util",
     ],
 )
 
@@ -129,6 +139,8 @@ multi_substrate_py_test(
         # numpy dep
         # tensorflow dep,
         "//tensorflow_probability",
+        "//tensorflow_probability/python/distributions:mvn_diag",
+        "//tensorflow_probability/python/internal:samplers",
         "//tensorflow_probability/python/internal:test_util",
     ],
 )

tensorflow_probability/python/experimental/bijectors/__init__.py

@@ -22,8 +22,8 @@
 from tensorflow_probability.python.experimental.bijectors.highway_flow import HighwayFlow
 __all__ = [
     'build_highway_flow_layer',
-    'HighwayFlow',
     'forward_log_det_jacobian_ratio',
+    'HighwayFlow',
     'inverse_log_det_jacobian_ratio',
     'make_distribution_bijector',
     'ScalarFunctionWithInferredInverse',

tensorflow_probability/python/experimental/bijectors/highway_flow_test.py

@@ -16,11 +16,10 @@
 import tensorflow.compat.v2 as tf
 
 import tensorflow_probability as tfp
+from tensorflow_probability.python.distributions import mvn_diag
 from tensorflow_probability.python.internal import samplers
 from tensorflow_probability.python.internal import test_util
 
-tfb = tfp.bijectors
-tfd = tfp.distributions
 
 @test_util.test_all_tf_execution_regimes
 class HighwayFlowTests(test_util.TestCase):
@@ -116,8 +115,8 @@ def testResidualFractionGradientsWithCenteredDifference(self):
       lower_diagonal_weights_matrix=tf.eye(width),
       gate_first_n=width
     )
-    target = tfd.MultivariateNormalDiag(loc=tf.zeros(width),
-                                        scale_diag=tf.ones(width))
+    target = mvn_diag.MultivariateNormalDiag(loc=tf.zeros(width),
+                                             scale_diag=tf.ones(width))
     x = tf.ones((batch_size, width))
     with tf.GradientTape() as g:
       g.watch(bijector.residual_fraction)