From a77f8dd6c3ad5f442e6886835192c92fb520b6ce Mon Sep 17 00:00:00 2001 From: thomaswc Date: Wed, 27 Mar 2024 13:04:49 -0700 Subject: [PATCH] Move autobnn to spinoffs. PiperOrigin-RevId: 619626204 --- .../experimental => spinoffs}/autobnn/BUILD | 10 +++++----- .../autobnn/README.md | 0 .../autobnn/__init__.py | 18 +++++++++--------- .../experimental => spinoffs}/autobnn/bnn.py | 2 +- .../autobnn/bnn_test.py | 2 +- .../autobnn/bnn_tree.py | 8 ++++---- .../autobnn/bnn_tree_test.py | 4 ++-- .../autobnn/estimators.py | 8 ++++---- .../autobnn/estimators_test.py | 8 ++++---- .../autobnn/kernels.py | 2 +- .../autobnn/kernels_test.py | 4 ++-- .../autobnn/likelihoods.py | 0 .../autobnn/likelihoods_test.py | 2 +- .../autobnn/models.py | 10 +++++----- .../autobnn/models_test.py | 4 ++-- .../autobnn/operators.py | 4 ++-- .../autobnn/operators_test.py | 6 +++--- .../autobnn/setup_autobnn.sh | 0 .../autobnn/training_util.py | 4 ++-- .../autobnn/training_util_test.py | 8 ++++---- .../experimental => spinoffs}/autobnn/util.py | 2 +- .../autobnn/util_test.py | 4 ++-- 22 files changed, 55 insertions(+), 55 deletions(-) rename tensorflow_probability/{python/experimental => spinoffs}/autobnn/BUILD (95%) rename tensorflow_probability/{python/experimental => spinoffs}/autobnn/README.md (100%) rename tensorflow_probability/{python/experimental => spinoffs}/autobnn/__init__.py (61%) rename tensorflow_probability/{python/experimental => spinoffs}/autobnn/bnn.py (98%) rename tensorflow_probability/{python/experimental => spinoffs}/autobnn/bnn_test.py (97%) rename tensorflow_probability/{python/experimental => spinoffs}/autobnn/bnn_tree.py (94%) rename tensorflow_probability/{python/experimental => spinoffs}/autobnn/bnn_tree_test.py (96%) rename tensorflow_probability/{python/experimental => spinoffs}/autobnn/estimators.py (97%) rename tensorflow_probability/{python/experimental => spinoffs}/autobnn/estimators_test.py (93%) rename tensorflow_probability/{python/experimental => spinoffs}/autobnn/kernels.py (99%) rename tensorflow_probability/{python/experimental => spinoffs}/autobnn/kernels_test.py (98%) rename tensorflow_probability/{python/experimental => spinoffs}/autobnn/likelihoods.py (100%) rename tensorflow_probability/{python/experimental => spinoffs}/autobnn/likelihoods_test.py (96%) rename tensorflow_probability/{python/experimental => spinoffs}/autobnn/models.py (96%) rename tensorflow_probability/{python/experimental => spinoffs}/autobnn/models_test.py (93%) rename tensorflow_probability/{python/experimental => spinoffs}/autobnn/operators.py (98%) rename tensorflow_probability/{python/experimental => spinoffs}/autobnn/operators_test.py (97%) rename tensorflow_probability/{python/experimental => spinoffs}/autobnn/setup_autobnn.sh (100%) rename tensorflow_probability/{python/experimental => spinoffs}/autobnn/training_util.py (98%) rename tensorflow_probability/{python/experimental => spinoffs}/autobnn/training_util_test.py (95%) rename tensorflow_probability/{python/experimental => spinoffs}/autobnn/util.py (97%) rename tensorflow_probability/{python/experimental => spinoffs}/autobnn/util_test.py (94%) diff --git a/tensorflow_probability/python/experimental/autobnn/BUILD b/tensorflow_probability/spinoffs/autobnn/BUILD similarity index 95% rename from tensorflow_probability/python/experimental/autobnn/BUILD rename to tensorflow_probability/spinoffs/autobnn/BUILD index 6cff94bc4c..0fac911f7c 100644 --- a/tensorflow_probability/python/experimental/autobnn/BUILD +++ b/tensorflow_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/tensorflow_probability/python/experimental/autobnn/README.md b/tensorflow_probability/spinoffs/autobnn/README.md similarity index 100% rename from tensorflow_probability/python/experimental/autobnn/README.md rename to tensorflow_probability/spinoffs/autobnn/README.md diff --git a/tensorflow_probability/python/experimental/autobnn/__init__.py b/tensorflow_probability/spinoffs/autobnn/__init__.py similarity index 61% rename from tensorflow_probability/python/experimental/autobnn/__init__.py rename to tensorflow_probability/spinoffs/autobnn/__init__.py index d140213e99..7d7e5492db 100644 --- a/tensorflow_probability/python/experimental/autobnn/__init__.py +++ b/tensorflow_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/tensorflow_probability/python/experimental/autobnn/bnn.py b/tensorflow_probability/spinoffs/autobnn/bnn.py similarity index 98% rename from tensorflow_probability/python/experimental/autobnn/bnn.py rename to tensorflow_probability/spinoffs/autobnn/bnn.py index d153d7d9ba..20d6695861 100644 --- a/tensorflow_probability/python/experimental/autobnn/bnn.py +++ b/tensorflow_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/tensorflow_probability/python/experimental/autobnn/bnn_test.py b/tensorflow_probability/spinoffs/autobnn/bnn_test.py similarity index 97% rename from tensorflow_probability/python/experimental/autobnn/bnn_test.py rename to tensorflow_probability/spinoffs/autobnn/bnn_test.py index 53bb7b09e4..4d0fe04831 100644 --- a/tensorflow_probability/python/experimental/autobnn/bnn_test.py +++ b/tensorflow_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/tensorflow_probability/python/experimental/autobnn/bnn_tree.py b/tensorflow_probability/spinoffs/autobnn/bnn_tree.py similarity index 94% rename from tensorflow_probability/python/experimental/autobnn/bnn_tree.py rename to tensorflow_probability/spinoffs/autobnn/bnn_tree.py index 4e63635d2a..644e128307 100644 --- a/tensorflow_probability/python/experimental/autobnn/bnn_tree.py +++ b/tensorflow_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/tensorflow_probability/python/experimental/autobnn/bnn_tree_test.py b/tensorflow_probability/spinoffs/autobnn/bnn_tree_test.py similarity index 96% rename from tensorflow_probability/python/experimental/autobnn/bnn_tree_test.py rename to tensorflow_probability/spinoffs/autobnn/bnn_tree_test.py index 197c535dbe..ca1c23d9cf 100644 --- a/tensorflow_probability/python/experimental/autobnn/bnn_tree_test.py +++ b/tensorflow_probability/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/tensorflow_probability/python/experimental/autobnn/estimators.py b/tensorflow_probability/spinoffs/autobnn/estimators.py similarity index 97% rename from tensorflow_probability/python/experimental/autobnn/estimators.py rename to tensorflow_probability/spinoffs/autobnn/estimators.py index 18f3fdfeb2..5002ed2cc5 100644 --- a/tensorflow_probability/python/experimental/autobnn/estimators.py +++ b/tensorflow_probability/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/tensorflow_probability/python/experimental/autobnn/estimators_test.py b/tensorflow_probability/spinoffs/autobnn/estimators_test.py similarity index 93% rename from tensorflow_probability/python/experimental/autobnn/estimators_test.py rename to tensorflow_probability/spinoffs/autobnn/estimators_test.py index c1d9d28ef3..0f5f6c2b44 100644 --- a/tensorflow_probability/python/experimental/autobnn/estimators_test.py +++ b/tensorflow_probability/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/tensorflow_probability/python/experimental/autobnn/kernels.py b/tensorflow_probability/spinoffs/autobnn/kernels.py similarity index 99% rename from tensorflow_probability/python/experimental/autobnn/kernels.py rename to tensorflow_probability/spinoffs/autobnn/kernels.py index cc8cfc4c23..3c9ad8019e 100644 --- a/tensorflow_probability/python/experimental/autobnn/kernels.py +++ b/tensorflow_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/tensorflow_probability/python/experimental/autobnn/kernels_test.py b/tensorflow_probability/spinoffs/autobnn/kernels_test.py similarity index 98% rename from tensorflow_probability/python/experimental/autobnn/kernels_test.py rename to tensorflow_probability/spinoffs/autobnn/kernels_test.py index c5f26882eb..96362e580b 100644 --- a/tensorflow_probability/python/experimental/autobnn/kernels_test.py +++ b/tensorflow_probability/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/tensorflow_probability/python/experimental/autobnn/likelihoods.py b/tensorflow_probability/spinoffs/autobnn/likelihoods.py similarity index 100% rename from tensorflow_probability/python/experimental/autobnn/likelihoods.py rename to tensorflow_probability/spinoffs/autobnn/likelihoods.py diff --git a/tensorflow_probability/python/experimental/autobnn/likelihoods_test.py b/tensorflow_probability/spinoffs/autobnn/likelihoods_test.py similarity index 96% rename from tensorflow_probability/python/experimental/autobnn/likelihoods_test.py rename to tensorflow_probability/spinoffs/autobnn/likelihoods_test.py index 4776f951a3..a36a415a68 100644 --- a/tensorflow_probability/python/experimental/autobnn/likelihoods_test.py +++ b/tensorflow_probability/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/tensorflow_probability/python/experimental/autobnn/models.py b/tensorflow_probability/spinoffs/autobnn/models.py similarity index 96% rename from tensorflow_probability/python/experimental/autobnn/models.py rename to tensorflow_probability/spinoffs/autobnn/models.py index 31d32f37bb..8eb77d849d 100644 --- a/tensorflow_probability/python/experimental/autobnn/models.py +++ b/tensorflow_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/tensorflow_probability/python/experimental/autobnn/models_test.py b/tensorflow_probability/spinoffs/autobnn/models_test.py similarity index 93% rename from tensorflow_probability/python/experimental/autobnn/models_test.py rename to tensorflow_probability/spinoffs/autobnn/models_test.py index ed9d44ca6d..68b3a40044 100644 --- a/tensorflow_probability/python/experimental/autobnn/models_test.py +++ b/tensorflow_probability/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/tensorflow_probability/python/experimental/autobnn/operators.py b/tensorflow_probability/spinoffs/autobnn/operators.py similarity index 98% rename from tensorflow_probability/python/experimental/autobnn/operators.py rename to tensorflow_probability/spinoffs/autobnn/operators.py index 2872bc46ec..43e37b3250 100644 --- a/tensorflow_probability/python/experimental/autobnn/operators.py +++ b/tensorflow_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/tensorflow_probability/python/experimental/autobnn/operators_test.py b/tensorflow_probability/spinoffs/autobnn/operators_test.py similarity index 97% rename from tensorflow_probability/python/experimental/autobnn/operators_test.py rename to tensorflow_probability/spinoffs/autobnn/operators_test.py index aab0c34e5b..d22ec6e85e 100644 --- a/tensorflow_probability/python/experimental/autobnn/operators_test.py +++ b/tensorflow_probability/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/tensorflow_probability/python/experimental/autobnn/setup_autobnn.sh b/tensorflow_probability/spinoffs/autobnn/setup_autobnn.sh similarity index 100% rename from tensorflow_probability/python/experimental/autobnn/setup_autobnn.sh rename to tensorflow_probability/spinoffs/autobnn/setup_autobnn.sh diff --git a/tensorflow_probability/python/experimental/autobnn/training_util.py b/tensorflow_probability/spinoffs/autobnn/training_util.py similarity index 98% rename from tensorflow_probability/python/experimental/autobnn/training_util.py rename to tensorflow_probability/spinoffs/autobnn/training_util.py index 60bce255bb..1167d157d5 100644 --- a/tensorflow_probability/python/experimental/autobnn/training_util.py +++ b/tensorflow_probability/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/tensorflow_probability/python/experimental/autobnn/training_util_test.py b/tensorflow_probability/spinoffs/autobnn/training_util_test.py similarity index 95% rename from tensorflow_probability/python/experimental/autobnn/training_util_test.py rename to tensorflow_probability/spinoffs/autobnn/training_util_test.py index b746a4f0cb..52d7762080 100644 --- a/tensorflow_probability/python/experimental/autobnn/training_util_test.py +++ b/tensorflow_probability/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/tensorflow_probability/python/experimental/autobnn/util.py b/tensorflow_probability/spinoffs/autobnn/util.py similarity index 97% rename from tensorflow_probability/python/experimental/autobnn/util.py rename to tensorflow_probability/spinoffs/autobnn/util.py index 737ff00088..10674126d9 100644 --- a/tensorflow_probability/python/experimental/autobnn/util.py +++ b/tensorflow_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/tensorflow_probability/python/experimental/autobnn/util_test.py b/tensorflow_probability/spinoffs/autobnn/util_test.py similarity index 94% rename from tensorflow_probability/python/experimental/autobnn/util_test.py rename to tensorflow_probability/spinoffs/autobnn/util_test.py index dbaa3866c9..250d5099e9 100644 --- a/tensorflow_probability/python/experimental/autobnn/util_test.py +++ b/tensorflow_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):