Move composite_tensor.py into tfp.experimental.util.

This is so I can add a tfp.experimental.math without triggering the usual
Python import shadowing issue.

PiperOrigin-RevId: 390180664

Author: SiegeLordEx

tensorflow_probability/python/experimental/BUILD

@@ -39,7 +39,6 @@ multi_substrate_py_library(
     ],
     srcs_version = "PY3",
     substrates_omit_deps = [
-        ":composite_tensor",
         "//tensorflow_probability/python/experimental/auto_batching",
         "//tensorflow_probability/python/experimental/linalg",
         "//tensorflow_probability/python/experimental/marginalize",
@@ -49,9 +48,9 @@ multi_substrate_py_library(
         "//tensorflow_probability/python/experimental/substrates",
         "//tensorflow_probability/python/experimental/vi",
         "//tensorflow_probability/python/internal:auto_composite_tensor",
+        "//tensorflow_probability/python/experimental/util:composite_tensor",
     ],
     deps = [
-        ":composite_tensor",
         "//tensorflow_probability/python/experimental/auto_batching",
         "//tensorflow_probability/python/experimental/bijectors",
         "//tensorflow_probability/python/experimental/distribute",
@@ -66,31 +65,9 @@ multi_substrate_py_library(
         "//tensorflow_probability/python/experimental/sts_gibbs",
         "//tensorflow_probability/python/experimental/substrates",
         "//tensorflow_probability/python/experimental/util",
+        "//tensorflow_probability/python/experimental/util:composite_tensor",
         "//tensorflow_probability/python/experimental/vi",
         "//tensorflow_probability/python/internal:all_util",
         "//tensorflow_probability/python/internal:auto_composite_tensor",
     ],
 )
-
-py_library(
-    name = "composite_tensor",
-    srcs = ["composite_tensor.py"],
-    srcs_version = "PY3",
-    deps = [
-        # tensorflow dep,
-        "//tensorflow_probability/python/distributions",
-    ],
-)
-
-py_test(
-    name = "composite_tensor_test",
-    srcs = ["composite_tensor_test.py"],
-    python_version = "PY3",
-    srcs_version = "PY3",
-    deps = [
-        # numpy dep,
-        # tensorflow dep,
-        "//tensorflow_probability",
-        "//tensorflow_probability/python/internal:test_util",
-    ],
-)

tensorflow_probability/python/experimental/__init__.py

@@ -45,8 +45,8 @@
 from tensorflow_probability.python.experimental import substrates
 from tensorflow_probability.python.experimental import util
 from tensorflow_probability.python.experimental import vi
-from tensorflow_probability.python.experimental.composite_tensor import as_composite
-from tensorflow_probability.python.experimental.composite_tensor import register_composite
+from tensorflow_probability.python.experimental.util.composite_tensor import as_composite
+from tensorflow_probability.python.experimental.util.composite_tensor import register_composite
 from tensorflow_probability.python.internal import all_util
 from tensorflow_probability.python.internal.auto_composite_tensor import auto_composite_tensor
 from tensorflow_probability.python.internal.auto_composite_tensor import AutoCompositeTensor

tensorflow_probability/python/experimental/util/BUILD

@@ -37,10 +37,12 @@ multi_substrate_py_library(
     numpy_omit_deps = [":jit_public_methods"],
     srcs_version = "PY3",
     substrates_omit_deps = [
+        ":composite_tensor",
         ":deferred_module",
         ":trainable",
     ],
     deps = [
+        ":composite_tensor",
         ":deferred_module",
         ":jit_public_methods",
         ":trainable",
@@ -133,3 +135,26 @@ py_test(
         "//tensorflow_probability/python/internal:test_util",
     ],
 )
+
+py_library(
+    name = "composite_tensor",
+    srcs = ["composite_tensor.py"],
+    srcs_version = "PY3",
+    deps = [
+        # tensorflow dep,
+        "//tensorflow_probability/python/distributions",
+    ],
+)
+
+py_test(
+    name = "composite_tensor_test",
+    srcs = ["composite_tensor_test.py"],
+    python_version = "PY3",
+    srcs_version = "PY3",
+    deps = [
+        # numpy dep,
+        # tensorflow dep,
+        "//tensorflow_probability",
+        "//tensorflow_probability/python/internal:test_util",
+    ],
+)

tensorflow_probability/python/experimental/util/__init__.py

@@ -18,6 +18,8 @@
 from __future__ import division
 from __future__ import print_function
 
+from tensorflow_probability.python.experimental.util.composite_tensor import as_composite
+from tensorflow_probability.python.experimental.util.composite_tensor import register_composite
 from tensorflow_probability.python.experimental.util.deferred_module import DeferredModule
 from tensorflow_probability.python.experimental.util.jit_public_methods import DEFAULT_METHODS_EXCLUDED_FROM_JIT
 from tensorflow_probability.python.experimental.util.jit_public_methods import JitPublicMethods
@@ -26,10 +28,12 @@
 
 
 _allowed_symbols = [
+    'as_composite',
     'DEFAULT_METHODS_EXCLUDED_FROM_JIT',
     'DeferredModule',
     'JitPublicMethods',
-    'make_trainable'
+    'make_trainable',
+    'register_composite',
 ]
 
 all_util.remove_undocumented(__name__, _allowed_symbols)

tensorflow_probability/python/experimental/composite_tensor.py renamed to tensorflow_probability/python/experimental/util/composite_tensor.py

@@ -26,7 +26,7 @@
 import tensorflow.compat.v2 as tf
 import tensorflow_probability as tfp
 
-from tensorflow_probability.python.experimental.composite_tensor import _registry as clsid_registry
+from tensorflow_probability.python.experimental.util.composite_tensor import _registry as clsid_registry
 from tensorflow_probability.python.internal import test_util as tfp_test_util
 from tensorflow.python.framework import test_util  # pylint: disable=g-direct-tensorflow-import