Remove tfp.distributions dependency from tfp.experimental.composite_tensor and import it only when needed (reduces integration test footprint).

PiperOrigin-RevId: 476165470

Author: emilyfertig

tensorflow_probability/python/experimental/util/BUILD

@@ -156,7 +156,6 @@ py_library(
     srcs = ["composite_tensor.py"],
     deps = [
         # tensorflow dep,
-        "//tensorflow_probability/python/distributions",
     ],
 )
 

tensorflow_probability/python/experimental/util/composite_tensor.py

@@ -18,7 +18,6 @@
 import six
 
 import tensorflow.compat.v2 as tf
-from tensorflow_probability.python import distributions
 from tensorflow_probability.python.internal import tensor_util
 from tensorflow.python.framework.composite_tensor import CompositeTensor  # pylint: disable=g-direct-tensorflow-import
 from tensorflow.python.saved_model import nested_structure_coder  # pylint: disable=g-direct-tensorflow-import
@@ -160,6 +159,7 @@ def _find_clsid(clsid):
   pkg, cls = clsid
   if clsid not in _registry:
     if pkg.startswith('tensorflow_probability.') and '.distributions' in pkg:
+      from tensorflow_probability.python import distributions  # pylint: disable=g-import-not-at-top
       dist_cls = getattr(distributions, cls)
       if (inspect.isclass(dist_cls) and
           issubclass(dist_cls, distributions.Distribution)):

tensorflow_probability/python/experimental/vi/util/BUILD

@@ -48,6 +48,7 @@ multi_substrate_py_library(
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:prefer_static",
         "//tensorflow_probability/python/internal:samplers",
+        "//tensorflow_probability/python/internal:trainable_state_util",
     ],
 )
 

tensorflow_probability/python/experimental/vi/util/trainable_linear_operators.py

@@ -218,7 +218,7 @@ def _trainable_linear_operator_tril(
     scale_tril_bijector = fill_scale_tril.FillScaleTriL(
         diag_bijector, diag_shift=tf.zeros([], dtype=dtype))
     scale_tril = yield trainable_state_util.Parameter(
-        init_fn=lambda seed: scale_tril_bijector(  # pylint: disable=g-long-lambda
+        init_fn=lambda seed: scale_tril_bijector.forward(  # pylint: disable=g-long-lambda
             samplers.normal(
                 mean=0.,
                 stddev=scale_initializer,
@@ -262,7 +262,7 @@ def _trainable_linear_operator_diag(
 
     diag_bijector = diag_bijector or _DefaultScaleDiagonal()
     scale_diag = yield trainable_state_util.Parameter(
-        init_fn=lambda seed: diag_bijector(  # pylint: disable=g-long-lambda
+        init_fn=lambda seed: diag_bijector.forward(  # pylint: disable=g-long-lambda
             samplers.normal(
                 mean=0.,
                 stddev=scale_initializer,