Make spinoffs/autobnn a pip installable package.

PiperOrigin-RevId: 619748042

tensorflow_probability/spinoffs/autobnn/BUILD → spinoffs/autobnn/autobnn/BUILD

@@ -37,7 +37,7 @@ py_library(
         ":operators",
         ":training_util",
         ":util",
-        "//tensorflow_probability/python/internal:all_util",
+        # tensorflow_probability/python/internal:all_util dep,
     ],
 )
 
@@ -49,7 +49,7 @@ py_library(
         # flax:core dep,
         # jax dep,
         # jaxtyping dep,
-        "//tensorflow_probability/python/distributions:distribution.jax",
+        # tensorflow_probability/python/distributions:distribution.jax dep,
     ],
 )
 
@@ -62,8 +62,8 @@ py_test(
         # google/protobuf:use_fast_cpp_protos dep,
         # jax dep,
         "//tensorflow_probability:jax",
-        "//tensorflow_probability/python/distributions:lognormal.jax",
-        "//tensorflow_probability/python/distributions:normal.jax",
+        # tensorflow_probability/python/distributions:lognormal.jax dep,
+        # tensorflow_probability/python/distributions:normal.jax dep,
     ],
 )
 
@@ -118,7 +118,7 @@ py_test(
         ":estimators",
         ":kernels",
         ":operators",
-        "//tensorflow_probability/python/internal:test_util",
+        # tensorflow_probability/python/internal:test_util dep,
     ],
 )
 
@@ -130,10 +130,10 @@ py_library(
         # flax dep,
         # flax:core dep,
         # jax dep,
-        "//tensorflow_probability/python/distributions:lognormal.jax",
-        "//tensorflow_probability/python/distributions:normal.jax",
-        "//tensorflow_probability/python/distributions:student_t.jax",
-        "//tensorflow_probability/python/distributions:uniform.jax",
+        # tensorflow_probability/python/distributions:lognormal.jax dep,
+        # tensorflow_probability/python/distributions:normal.jax dep,
+        # tensorflow_probability/python/distributions:student_t.jax dep,
+        # tensorflow_probability/python/distributions:uniform.jax dep,
     ],
 )
 
@@ -147,7 +147,7 @@ py_test(
         # absl/testing:parameterized dep,
         # google/protobuf:use_fast_cpp_protos dep,
         # jax dep,
-        "//tensorflow_probability/python/distributions:lognormal.jax",
+        # tensorflow_probability/python/distributions:lognormal.jax dep,
     ],
 )
 
@@ -158,14 +158,14 @@ py_library(
         # flax:core dep,
         # jax dep,
         # jaxtyping dep,
-        "//tensorflow_probability/python/bijectors:softplus.jax",
-        "//tensorflow_probability/python/distributions:distribution.jax",
-        "//tensorflow_probability/python/distributions:inflated.jax",
-        "//tensorflow_probability/python/distributions:logistic.jax",
-        "//tensorflow_probability/python/distributions:lognormal.jax",
-        "//tensorflow_probability/python/distributions:negative_binomial.jax",
-        "//tensorflow_probability/python/distributions:normal.jax",
-        "//tensorflow_probability/python/distributions:transformed_distribution.jax",
+        # tensorflow_probability/python/bijectors:softplus.jax dep,
+        # tensorflow_probability/python/distributions:distribution.jax dep,
+        # tensorflow_probability/python/distributions:inflated.jax dep,
+        # tensorflow_probability/python/distributions:logistic.jax dep,
+        # tensorflow_probability/python/distributions:lognormal.jax dep,
+        # tensorflow_probability/python/distributions:negative_binomial.jax dep,
+        # tensorflow_probability/python/distributions:normal.jax dep,
+        # tensorflow_probability/python/distributions:transformed_distribution.jax dep,
     ],
 )
 
@@ -216,14 +216,14 @@ py_library(
         ":likelihoods",
         # flax:core dep,
         # jax dep,
-        "//tensorflow_probability/python/bijectors:chain.jax",
-        "//tensorflow_probability/python/bijectors:scale.jax",
-        "//tensorflow_probability/python/bijectors:shift.jax",
-        "//tensorflow_probability/python/distributions:beta.jax",
-        "//tensorflow_probability/python/distributions:dirichlet.jax",
-        "//tensorflow_probability/python/distributions:half_normal.jax",
-        "//tensorflow_probability/python/distributions:normal.jax",
-        "//tensorflow_probability/python/distributions:transformed_distribution.jax",
+        # tensorflow_probability/python/bijectors:chain.jax dep,
+        # tensorflow_probability/python/bijectors:scale.jax dep,
+        # tensorflow_probability/python/bijectors:shift.jax dep,
+        # tensorflow_probability/python/distributions:beta.jax dep,
+        # tensorflow_probability/python/distributions:dirichlet.jax dep,
+        # tensorflow_probability/python/distributions:half_normal.jax dep,
+        # tensorflow_probability/python/distributions:normal.jax dep,
+        # tensorflow_probability/python/distributions:transformed_distribution.jax dep,
     ],
 )
 
@@ -242,7 +242,7 @@ py_test(
         # google/protobuf:use_fast_cpp_protos dep,
         # jax dep,
         # numpy dep,
-        "//tensorflow_probability/python/distributions:distribution.jax",
+        # tensorflow_probability/python/distributions:distribution.jax dep,
     ],
 )
 
@@ -258,7 +258,6 @@ py_library(
         # matplotlib dep,
         # numpy dep,
         # pandas dep,
-        "//tensorflow_probability/python/experimental/timeseries:metrics",
     ],
 )
 
@@ -276,7 +275,7 @@ py_test(
         # google/protobuf:use_fast_cpp_protos dep,
         # jax dep,
         # numpy dep,
-        "//tensorflow_probability/python/internal:test_util",
+        # tensorflow_probability/python/internal:test_util dep,
     ],
 )
 
@@ -288,7 +287,7 @@ py_library(
         # jax dep,
         # numpy dep,
         # scipy dep,
-        "//tensorflow_probability/python/distributions:distribution.jax",
+        # tensorflow_probability/python/distributions:distribution.jax dep,
     ],
 )
 
@@ -301,6 +300,6 @@ py_test(
         # google/protobuf:use_fast_cpp_protos dep,
         # jax dep,
         # numpy dep,
-        "//tensorflow_probability/python/internal:test_util",
+        # tensorflow_probability/python/internal:test_util dep,
     ],
 )

tensorflow_probability/spinoffs/autobnn/__init__.py → spinoffs/autobnn/autobnn/__init__.py

@@ -14,18 +14,17 @@
 # ============================================================================
 """Package for training GP-like Bayesian Neural Nets w/ composite structure."""
 
-from tensorflow_probability.python.internal import all_util
-from tensorflow_probability.spinoffs.autobnn import bnn
-from tensorflow_probability.spinoffs.autobnn import bnn_tree
-from tensorflow_probability.spinoffs.autobnn import estimators
-from tensorflow_probability.spinoffs.autobnn import kernels
-from tensorflow_probability.spinoffs.autobnn import likelihoods
-from tensorflow_probability.spinoffs.autobnn import models
-from tensorflow_probability.spinoffs.autobnn import operators
-from tensorflow_probability.spinoffs.autobnn import training_util
-from tensorflow_probability.spinoffs.autobnn import util
+from autobnn import bnn
+from autobnn import bnn_tree
+from autobnn import estimators
+from autobnn import kernels
+from autobnn import likelihoods
+from autobnn import models
+from autobnn import operators
+from autobnn import training_util
+from autobnn import util
 
-_allowed_symbols = [
+__all__ = [
     'bnn',
     'bnn_tree',
     'estimators',
@@ -36,5 +35,3 @@
     'training_util',
     'util',
 ]
-
-all_util.remove_undocumented(__name__, _allowed_symbols)

tensorflow_probability/spinoffs/autobnn/bnn.py → spinoffs/autobnn/autobnn/bnn.py

@@ -21,7 +21,7 @@
 import jax
 import jax.numpy as jnp
 from jaxtyping import Array, Float, PyTree  # pylint: disable=g-importing-member,g-multiple-import
-from tensorflow_probability.spinoffs.autobnn import likelihoods
+from autobnn import likelihoods
 from tensorflow_probability.substrates.jax.distributions import distribution as distribution_lib
 
 

tensorflow_probability/spinoffs/autobnn/bnn_test.py → spinoffs/autobnn/autobnn/bnn_test.py

@@ -17,7 +17,7 @@
 from flax import linen as nn
 import jax
 import jax.numpy as jnp
-from tensorflow_probability.spinoffs.autobnn import bnn
+from autobnn import bnn
 from tensorflow_probability.substrates.jax.distributions import lognormal as lognormal_lib
 from tensorflow_probability.substrates.jax.distributions import normal as normal_lib
 from absl.testing import absltest

tensorflow_probability/spinoffs/autobnn/bnn_tree.py → spinoffs/autobnn/autobnn/bnn_tree.py

@@ -19,10 +19,10 @@
 from flax import linen as nn
 import jax
 import jax.numpy as jnp
-from tensorflow_probability.spinoffs.autobnn import bnn
-from tensorflow_probability.spinoffs.autobnn import kernels
-from tensorflow_probability.spinoffs.autobnn import operators
-from tensorflow_probability.spinoffs.autobnn import util
+from autobnn import bnn
+from autobnn import kernels
+from autobnn import operators
+from autobnn import util
 
 Array = jnp.ndarray
 

tensorflow_probability/spinoffs/autobnn/bnn_tree_test.py → spinoffs/autobnn/autobnn/bnn_tree_test.py

@@ -18,8 +18,8 @@
 from flax import linen as nn
 import jax
 import jax.numpy as jnp
-from tensorflow_probability.spinoffs.autobnn import bnn_tree
-from tensorflow_probability.spinoffs.autobnn import kernels
+from autobnn import bnn_tree
+from autobnn import kernels
 from absl.testing import absltest
 
 

tensorflow_probability/spinoffs/autobnn/estimators.py → spinoffs/autobnn/autobnn/estimators.py

@@ -19,10 +19,10 @@
 import jax
 import jax.numpy as jnp
 from jaxtyping import ArrayLike, PyTree  # pylint: disable=g-importing-member,g-multiple-import
-from tensorflow_probability.spinoffs.autobnn import bnn
-from tensorflow_probability.spinoffs.autobnn import likelihoods
-from tensorflow_probability.spinoffs.autobnn import models
-from tensorflow_probability.spinoffs.autobnn import training_util
+from autobnn import bnn
+from autobnn import likelihoods
+from autobnn import models
+from autobnn import training_util
 
 
 class _AutoBnnEstimator:

tensorflow_probability/spinoffs/autobnn/estimators_test.py → spinoffs/autobnn/autobnn/estimators_test.py

@@ -17,10 +17,10 @@
 import jax
 import numpy as np
 from tensorflow_probability.python.internal import test_util
-from tensorflow_probability.spinoffs.autobnn import estimators
-from tensorflow_probability.spinoffs.autobnn import kernels
-from tensorflow_probability.spinoffs.autobnn import operators
-from tensorflow_probability.spinoffs.autobnn import util
+from autobnn import estimators
+from autobnn import kernels
+from autobnn import operators
+from autobnn import util
 
 
 class AutoBNNTest(test_util.TestCase):

tensorflow_probability/spinoffs/autobnn/kernels.py → spinoffs/autobnn/autobnn/kernels.py

@@ -19,7 +19,7 @@
 from flax.linen import initializers
 import jax
 import jax.numpy as jnp
-from tensorflow_probability.spinoffs.autobnn import bnn
+from autobnn import bnn
 from tensorflow_probability.substrates.jax.distributions import lognormal as lognormal_lib
 from tensorflow_probability.substrates.jax.distributions import normal as normal_lib
 from tensorflow_probability.substrates.jax.distributions import student_t as student_t_lib

tensorflow_probability/spinoffs/autobnn/kernels_test.py → spinoffs/autobnn/autobnn/kernels_test.py

@@ -18,8 +18,8 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
-from tensorflow_probability.spinoffs.autobnn import kernels
-from tensorflow_probability.spinoffs.autobnn import util
+from autobnn import kernels
+from autobnn import util
 from tensorflow_probability.substrates.jax.distributions import lognormal as lognormal_lib
 
 from absl.testing import absltest

tensorflow_probability/spinoffs/autobnn/likelihoods_test.py → spinoffs/autobnn/autobnn/likelihoods_test.py

@@ -16,7 +16,7 @@
 
 from absl.testing import parameterized
 import jax.numpy as jnp
-from tensorflow_probability.spinoffs.autobnn import likelihoods
+from autobnn import likelihoods
 from absl.testing import absltest
 
 

tensorflow_probability/spinoffs/autobnn/models.py → spinoffs/autobnn/autobnn/models.py

@@ -22,11 +22,11 @@
 import functools
 from typing import Sequence, Union
 import jax.numpy as jnp
-from tensorflow_probability.spinoffs.autobnn import bnn
-from tensorflow_probability.spinoffs.autobnn import bnn_tree
-from tensorflow_probability.spinoffs.autobnn import kernels
-from tensorflow_probability.spinoffs.autobnn import likelihoods
-from tensorflow_probability.spinoffs.autobnn import operators
+from autobnn import bnn
+from autobnn import bnn_tree
+from autobnn import kernels
+from autobnn import likelihoods
+from autobnn import operators
 
 
 Array = jnp.ndarray

tensorflow_probability/spinoffs/autobnn/models_test.py → spinoffs/autobnn/autobnn/models_test.py

@@ -17,8 +17,8 @@
 from absl.testing import parameterized
 import jax
 import jax.numpy as jnp
-from tensorflow_probability.spinoffs.autobnn import likelihoods
-from tensorflow_probability.spinoffs.autobnn import models
+from autobnn import likelihoods
+from autobnn import models
 from absl.testing import absltest
 
 

tensorflow_probability/spinoffs/autobnn/operators.py → spinoffs/autobnn/autobnn/operators.py

@@ -19,8 +19,8 @@
 from flax import linen as nn
 import jax
 import jax.numpy as jnp
-from tensorflow_probability.spinoffs.autobnn import bnn
-from tensorflow_probability.spinoffs.autobnn import likelihoods
+from autobnn import bnn
+from autobnn import likelihoods
 from tensorflow_probability.substrates.jax.bijectors import chain as chain_lib
 from tensorflow_probability.substrates.jax.bijectors import scale as scale_lib
 from tensorflow_probability.substrates.jax.bijectors import shift as shift_lib

tensorflow_probability/spinoffs/autobnn/operators_test.py → spinoffs/autobnn/autobnn/operators_test.py

@@ -19,9 +19,9 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
-from tensorflow_probability.spinoffs.autobnn import kernels
-from tensorflow_probability.spinoffs.autobnn import operators
-from tensorflow_probability.spinoffs.autobnn import util
+from autobnn import kernels
+from autobnn import operators
+from autobnn import util
 from tensorflow_probability.substrates.jax.distributions import distribution as distribution_lib
 from absl.testing import absltest