Move autobnn to spinoffs.

PiperOrigin-RevId: 619626204
tensorflow · Mar 27, 2024 · a77f8dd · a77f8dd
1 parent 7e1ce86
commit a77f8dd
Show file tree

Hide file tree

Showing 22 changed files with 55 additions and 55 deletions.
diff --git a/...ability/python/experimental/autobnn/BUILD → ...orflow_probability/spinoffs/autobnn/BUILD b/...ability/python/experimental/autobnn/BUILD → ...orflow_probability/spinoffs/autobnn/BUILD
@@ -250,14 +250,14 @@ py_library(
     name = "training_util",
     srcs = ["training_util.py"],
     deps = [
+        ":bnn",
+        ":util",
         # bayeux dep,
         # jax dep,
         # jaxtyping dep,
         # matplotlib dep,
         # numpy dep,
         # pandas dep,
-        "//tensorflow_probability/python/experimental/autobnn:bnn",
-        "//tensorflow_probability/python/experimental/autobnn:util",
         "//tensorflow_probability/python/experimental/timeseries:metrics",
     ],
 )
@@ -268,14 +268,14 @@ py_test(
     # TODO(b/322864412): enable this test in OSS.
     tags = ["no-oss-ci"],
     deps = [
+        ":kernels",
+        ":operators",
         ":training_util",
+        ":util",
         # chex dep,
         # google/protobuf:use_fast_cpp_protos dep,
         # jax dep,
         # numpy dep,
-        "//tensorflow_probability/python/experimental/autobnn:kernels",
-        "//tensorflow_probability/python/experimental/autobnn:operators",
-        "//tensorflow_probability/python/experimental/autobnn:util",
         "//tensorflow_probability/python/internal:test_util",
     ],
 )

diff --git a/...ity/python/experimental/autobnn/README.md → ...ow_probability/spinoffs/autobnn/README.md b/...ity/python/experimental/autobnn/README.md → ...ow_probability/spinoffs/autobnn/README.md
diff --git a/...y/python/experimental/autobnn/__init__.py → ..._probability/spinoffs/autobnn/__init__.py b/...y/python/experimental/autobnn/__init__.py → ..._probability/spinoffs/autobnn/__init__.py
@@ -14,16 +14,16 @@
 # ============================================================================
 """Package for training GP-like Bayesian Neural Nets w/ composite structure."""
 
-from tensorflow_probability.python.experimental.autobnn import bnn
-from tensorflow_probability.python.experimental.autobnn import bnn_tree
-from tensorflow_probability.python.experimental.autobnn import estimators
-from tensorflow_probability.python.experimental.autobnn import kernels
-from tensorflow_probability.python.experimental.autobnn import likelihoods
-from tensorflow_probability.python.experimental.autobnn import models
-from tensorflow_probability.python.experimental.autobnn import operators
-from tensorflow_probability.python.experimental.autobnn import training_util
-from tensorflow_probability.python.experimental.autobnn import util
 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
 
 _allowed_symbols = [
     'bnn',

diff --git a/...bility/python/experimental/autobnn/bnn.py → ...rflow_probability/spinoffs/autobnn/bnn.py b/...bility/python/experimental/autobnn/bnn.py → ...rflow_probability/spinoffs/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.python.experimental.autobnn import likelihoods
+from tensorflow_probability.spinoffs.autobnn import likelihoods
 from tensorflow_probability.substrates.jax.distributions import distribution as distribution_lib
 
 

diff --git a/...y/python/experimental/autobnn/bnn_test.py → ..._probability/spinoffs/autobnn/bnn_test.py b/...y/python/experimental/autobnn/bnn_test.py → ..._probability/spinoffs/autobnn/bnn_test.py
@@ -17,7 +17,7 @@
 from flax import linen as nn
 import jax
 import jax.numpy as jnp
-from tensorflow_probability.python.experimental.autobnn import bnn
+from tensorflow_probability.spinoffs.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

diff --git a/...y/python/experimental/autobnn/bnn_tree.py → ..._probability/spinoffs/autobnn/bnn_tree.py b/...y/python/experimental/autobnn/bnn_tree.py → ..._probability/spinoffs/autobnn/bnn_tree.py
@@ -19,10 +19,10 @@
 from flax import linen as nn
 import jax
 import jax.numpy as jnp
-from tensorflow_probability.python.experimental.autobnn import bnn
-from tensorflow_probability.python.experimental.autobnn import kernels
-from tensorflow_probability.python.experimental.autobnn import operators
-from tensorflow_probability.python.experimental.autobnn import util
+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
 
 Array = jnp.ndarray
 

diff --git a/...hon/experimental/autobnn/bnn_tree_test.py → ...ability/spinoffs/autobnn/bnn_tree_test.py b/...hon/experimental/autobnn/bnn_tree_test.py → ...ability/spinoffs/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.python.experimental.autobnn import bnn_tree
-from tensorflow_probability.python.experimental.autobnn import kernels
+from tensorflow_probability.spinoffs.autobnn import bnn_tree
+from tensorflow_probability.spinoffs.autobnn import kernels
 from absl.testing import absltest
 
 

diff --git a/...python/experimental/autobnn/estimators.py → ...robability/spinoffs/autobnn/estimators.py b/...python/experimental/autobnn/estimators.py → ...robability/spinoffs/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.python.experimental.autobnn import bnn
-from tensorflow_probability.python.experimental.autobnn import likelihoods
-from tensorflow_probability.python.experimental.autobnn import models
-from tensorflow_probability.python.experimental.autobnn import training_util
+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
 
 
 class _AutoBnnEstimator:

diff --git a/...n/experimental/autobnn/estimators_test.py → ...ility/spinoffs/autobnn/estimators_test.py b/...n/experimental/autobnn/estimators_test.py → ...ility/spinoffs/autobnn/estimators_test.py
@@ -16,11 +16,11 @@
 
 import jax
 import numpy as np
-from tensorflow_probability.python.experimental.autobnn import estimators
-from tensorflow_probability.python.experimental.autobnn import kernels
-from tensorflow_probability.python.experimental.autobnn import operators
-from tensorflow_probability.python.experimental.autobnn import util
 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
 
 
 class AutoBNNTest(test_util.TestCase):

diff --git a/...ty/python/experimental/autobnn/kernels.py → ...w_probability/spinoffs/autobnn/kernels.py b/...ty/python/experimental/autobnn/kernels.py → ...w_probability/spinoffs/autobnn/kernels.py
@@ -19,7 +19,7 @@
 from flax.linen import initializers
 import jax
 import jax.numpy as jnp
-from tensorflow_probability.python.experimental.autobnn import bnn
+from tensorflow_probability.spinoffs.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

diff --git a/...thon/experimental/autobnn/kernels_test.py → ...bability/spinoffs/autobnn/kernels_test.py b/...thon/experimental/autobnn/kernels_test.py → ...bability/spinoffs/autobnn/kernels_test.py
@@ -18,8 +18,8 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
-from tensorflow_probability.python.experimental.autobnn import kernels
-from tensorflow_probability.python.experimental.autobnn import util
+from tensorflow_probability.spinoffs.autobnn import kernels
+from tensorflow_probability.spinoffs.autobnn import util
 from tensorflow_probability.substrates.jax.distributions import lognormal as lognormal_lib
 
 from absl.testing import absltest

diff --git a/...ython/experimental/autobnn/likelihoods.py → ...obability/spinoffs/autobnn/likelihoods.py b/...ython/experimental/autobnn/likelihoods.py → ...obability/spinoffs/autobnn/likelihoods.py
diff --git a/.../experimental/autobnn/likelihoods_test.py → ...lity/spinoffs/autobnn/likelihoods_test.py b/.../experimental/autobnn/likelihoods_test.py → ...lity/spinoffs/autobnn/likelihoods_test.py
@@ -16,7 +16,7 @@
 
 from absl.testing import parameterized
 import jax.numpy as jnp
-from tensorflow_probability.python.experimental.autobnn import likelihoods
+from tensorflow_probability.spinoffs.autobnn import likelihoods
 from absl.testing import absltest
 
 

diff --git a/...ity/python/experimental/autobnn/models.py → ...ow_probability/spinoffs/autobnn/models.py b/...ity/python/experimental/autobnn/models.py → ...ow_probability/spinoffs/autobnn/models.py
@@ -22,11 +22,11 @@
 import functools
 from typing import Sequence, Union
 import jax.numpy as jnp
-from tensorflow_probability.python.experimental.autobnn import bnn
-from tensorflow_probability.python.experimental.autobnn import bnn_tree
-from tensorflow_probability.python.experimental.autobnn import kernels
-from tensorflow_probability.python.experimental.autobnn import likelihoods
-from tensorflow_probability.python.experimental.autobnn import operators
+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
 
 
 Array = jnp.ndarray

diff --git a/...ython/experimental/autobnn/models_test.py → ...obability/spinoffs/autobnn/models_test.py b/...ython/experimental/autobnn/models_test.py → ...obability/spinoffs/autobnn/models_test.py
@@ -17,8 +17,8 @@
 from absl.testing import parameterized
 import jax
 import jax.numpy as jnp
-from tensorflow_probability.python.experimental.autobnn import likelihoods
-from tensorflow_probability.python.experimental.autobnn import models
+from tensorflow_probability.spinoffs.autobnn import likelihoods
+from tensorflow_probability.spinoffs.autobnn import models
 from absl.testing import absltest
 
 

diff --git a/.../python/experimental/autobnn/operators.py → ...probability/spinoffs/autobnn/operators.py b/.../python/experimental/autobnn/operators.py → ...probability/spinoffs/autobnn/operators.py
@@ -19,8 +19,8 @@
 from flax import linen as nn
 import jax
 import jax.numpy as jnp
-from tensorflow_probability.python.experimental.autobnn import bnn
-from tensorflow_probability.python.experimental.autobnn import likelihoods
+from tensorflow_probability.spinoffs.autobnn import bnn
+from tensorflow_probability.spinoffs.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

diff --git a/...on/experimental/autobnn/operators_test.py → ...bility/spinoffs/autobnn/operators_test.py b/...on/experimental/autobnn/operators_test.py → ...bility/spinoffs/autobnn/operators_test.py
@@ -19,9 +19,9 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
-from tensorflow_probability.python.experimental.autobnn import kernels
-from tensorflow_probability.python.experimental.autobnn import operators
-from tensorflow_probability.python.experimental.autobnn import util
+from tensorflow_probability.spinoffs.autobnn import kernels
+from tensorflow_probability.spinoffs.autobnn import operators
+from tensorflow_probability.spinoffs.autobnn import util
 from tensorflow_probability.substrates.jax.distributions import distribution as distribution_lib
 from absl.testing import absltest
 

diff --git a/...hon/experimental/autobnn/setup_autobnn.sh → ...ability/spinoffs/autobnn/setup_autobnn.sh b/...hon/experimental/autobnn/setup_autobnn.sh → ...ability/spinoffs/autobnn/setup_autobnn.sh
diff --git a/...hon/experimental/autobnn/training_util.py → ...ability/spinoffs/autobnn/training_util.py b/...hon/experimental/autobnn/training_util.py → ...ability/spinoffs/autobnn/training_util.py
@@ -23,9 +23,9 @@
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
-from tensorflow_probability.python.experimental.autobnn import bnn
-from tensorflow_probability.python.experimental.autobnn import util
 from tensorflow_probability.python.experimental.timeseries import metrics
+from tensorflow_probability.spinoffs.autobnn import bnn
+from tensorflow_probability.spinoffs.autobnn import util
 
 
 def _make_bayeux_model(

diff --git a/...xperimental/autobnn/training_util_test.py → ...ty/spinoffs/autobnn/training_util_test.py b/...xperimental/autobnn/training_util_test.py → ...ty/spinoffs/autobnn/training_util_test.py
@@ -18,11 +18,11 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
-from tensorflow_probability.python.experimental.autobnn import kernels
-from tensorflow_probability.python.experimental.autobnn import operators
-from tensorflow_probability.python.experimental.autobnn import training_util
-from tensorflow_probability.python.experimental.autobnn import util
 from tensorflow_probability.python.internal import test_util
+from tensorflow_probability.spinoffs.autobnn import kernels
+from tensorflow_probability.spinoffs.autobnn import operators
+from tensorflow_probability.spinoffs.autobnn import training_util
+from tensorflow_probability.spinoffs.autobnn import util
 
 
 class TrainingUtilTest(test_util.TestCase):

diff --git a/...ility/python/experimental/autobnn/util.py → ...flow_probability/spinoffs/autobnn/util.py b/...ility/python/experimental/autobnn/util.py → ...flow_probability/spinoffs/autobnn/util.py
@@ -19,7 +19,7 @@
 import jax
 import jax.numpy as jnp
 import scipy
-from tensorflow_probability.python.experimental.autobnn import bnn
+from tensorflow_probability.spinoffs.autobnn import bnn
 from tensorflow_probability.substrates.jax.distributions import distribution as distribution_lib
 
 

diff --git a/.../python/experimental/autobnn/util_test.py → ...probability/spinoffs/autobnn/util_test.py b/.../python/experimental/autobnn/util_test.py → ...probability/spinoffs/autobnn/util_test.py
@@ -17,9 +17,9 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
-from tensorflow_probability.python.experimental.autobnn import kernels
-from tensorflow_probability.python.experimental.autobnn import util
 from tensorflow_probability.python.internal import test_util
+from tensorflow_probability.spinoffs.autobnn import kernels
+from tensorflow_probability.spinoffs.autobnn import util
 
 
 class UtilTest(test_util.TestCase):