Remove fully_bayesian re-exports from botorch.models.__init__ (#3291)

Summary:
Pull Request resolved: #3291

X-link: facebook/Ax#5190

Differential Revision: D102367257

Author: sdaulton

botorch/__init__.py

@@ -16,7 +16,7 @@
     test_functions,
 )
 from botorch.cross_validation import batch_cross_validation
-from botorch.fit import fit_fully_bayesian_model_nuts, fit_gpytorch_mll
+from botorch.fit import fit_gpytorch_mll
 from botorch.generation.gen import (
     gen_candidates_scipy,
     gen_candidates_torch,
@@ -53,7 +53,6 @@
     "acquisition",
     "batch_cross_validation",
     "exceptions",
-    "fit_fully_bayesian_model_nuts",
     "fit_gpytorch_mll",
     "gen_candidates_scipy",
     "gen_candidates_torch",

botorch/fit.py

@@ -12,16 +12,17 @@
 from copy import deepcopy
 from functools import partial
 from itertools import filterfalse
-from typing import Any
+from typing import Any, TYPE_CHECKING
 from warnings import catch_warnings, simplefilter, warn_explicit, WarningMessage
 
-import jax
+if TYPE_CHECKING:
+    from botorch.models.fully_bayesian import AbstractFullyBayesianSingleTaskGP
+    from botorch.models.fully_bayesian_multitask import SaasFullyBayesianMultiTaskGP
+
 from botorch.exceptions.errors import ModelFittingError, UnsupportedError
 from botorch.exceptions.warnings import OptimizationWarning
 from botorch.logging import logger
 from botorch.models import SingleTaskGP
-from botorch.models.fully_bayesian import AbstractFullyBayesianSingleTaskGP
-from botorch.models.fully_bayesian_multitask import SaasFullyBayesianMultiTaskGP
 from botorch.models.map_saas import get_map_saas_model
 from botorch.models.model_list_gp_regression import ModelListGP
 from botorch.models.transforms.input import InputTransform
@@ -44,7 +45,6 @@
 from gpytorch.mlls.marginal_log_likelihood import MarginalLogLikelihood
 from gpytorch.mlls.sum_marginal_log_likelihood import SumMarginalLogLikelihood
 from linear_operator.utils.errors import NotPSDError
-from numpyro.infer import MCMC, NUTS
 from torch import device, Tensor
 from torch.nn import Parameter
 from torch.utils.data import DataLoader
@@ -367,6 +367,11 @@ def fit_fully_bayesian_model_nuts(
         >>> gp = SaasFullyBayesianSingleTaskGP(train_X, train_Y)
         >>> fit_fully_bayesian_model_nuts(gp)
     """
+    # Local import to avoid pulling in JAX/numpyro at module level,
+    # which would break environments without NumPy >= 2.0.
+    import jax
+    from numpyro.infer import MCMC, NUTS
+
     model.train()
 
     # Do inference with NUTS

botorch/models/__init__.py

@@ -14,8 +14,6 @@
     GenericDeterministicModel,
     PosteriorMeanModel,
 )
-from botorch.models.fully_bayesian import SaasFullyBayesianSingleTaskGP
-from botorch.models.fully_bayesian_multitask import SaasFullyBayesianMultiTaskGP
 from botorch.models.gp_regression import SingleTaskGP
 from botorch.models.gp_regression_fidelity import SingleTaskMultiFidelityGP
 from botorch.models.gp_regression_mixed import MixedSingleTaskGP
@@ -47,8 +45,6 @@
     "PairwiseGP",
     "PairwiseLaplaceMarginalLogLikelihood",
     "PosteriorMeanModel",
-    "SaasFullyBayesianMultiTaskGP",
-    "SaasFullyBayesianSingleTaskGP",
     "SingleTaskGP",
     "SingleTaskMultiFidelityGP",
     "SingleTaskVariationalGP",

botorch/posteriors/__init__.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-from botorch.posteriors.fully_bayesian import GaussianMixturePosterior
 from botorch.posteriors.gpytorch import GPyTorchPosterior
 from botorch.posteriors.higher_order import HigherOrderGPPosterior
 from botorch.posteriors.multitask import MultitaskGPPosterior
@@ -14,7 +13,6 @@
 from botorch.posteriors.transformed import TransformedPosterior
 
 __all__ = [
-    "GaussianMixturePosterior",
     "GPyTorchPosterior",
     "HigherOrderGPPosterior",
     "MultitaskGPPosterior",

test/acquisition/test_input_constructors.py

@@ -111,8 +111,9 @@
     project_to_target_fidelity,
 )
 from botorch.exceptions.errors import BotorchError, UnsupportedError
-from botorch.models import MultiTaskGP, SaasFullyBayesianSingleTaskGP, SingleTaskGP
+from botorch.models import MultiTaskGP, SingleTaskGP
 from botorch.models.deterministic import FixedSingleSampleModel
+from botorch.models.fully_bayesian import SaasFullyBayesianSingleTaskGP
 from botorch.models.model_list_gp_regression import ModelListGP
 from botorch.optim.optimize import optimize_acqf
 from botorch.sampling.normal import IIDNormalSampler, SobolQMCNormalSampler

test/models/test_fully_bayesian.py

@@ -12,7 +12,7 @@
 import jax.numpy as jnp
 import numpyro
 import torch
-from botorch import fit_fully_bayesian_model_nuts, utils
+from botorch import utils
 from botorch.acquisition.analytic import (
     ExpectedImprovement,
     PosteriorMean,
@@ -41,6 +41,7 @@
     qLogNoisyExpectedHypervolumeImprovement,
 )
 from botorch.acquisition.utils import prune_inferior_points
+from botorch.fit import fit_fully_bayesian_model_nuts
 from botorch.models import ModelList, ModelListGP
 from botorch.models.deterministic import GenericDeterministicModel
 from botorch.models.fully_bayesian import (

test/models/test_fully_bayesian_multitask.py

@@ -10,7 +10,6 @@
 import jax.numpy as jnp
 import numpyro.handlers
 import torch
-from botorch import fit_fully_bayesian_model_nuts
 from botorch.acquisition.analytic import (
     LogExpectedImprovement,
     PosteriorMean,
@@ -30,6 +29,7 @@
     qLogExpectedHypervolumeImprovement,
     qLogNoisyExpectedHypervolumeImprovement,
 )
+from botorch.fit import fit_fully_bayesian_model_nuts
 from botorch.models import ModelList, ModelListGP
 from botorch.models.deterministic import GenericDeterministicModel
 from botorch.models.fully_bayesian import (
@@ -48,7 +48,7 @@
 )
 from botorch.models.transforms.input import Normalize
 from botorch.models.transforms.outcome import Standardize
-from botorch.posteriors import GaussianMixturePosterior
+from botorch.posteriors.fully_bayesian import GaussianMixturePosterior
 from botorch.sampling.get_sampler import get_sampler
 from botorch.sampling.normal import IIDNormalSampler
 from botorch.utils.multi_objective.box_decompositions.non_dominated import (

test/utils/test_transforms.py

@@ -14,9 +14,9 @@
     HigherOrderGP,
     ModelList,
     PairwiseGP,
-    SaasFullyBayesianSingleTaskGP,
     SingleTaskGP,
 )
+from botorch.models.fully_bayesian import SaasFullyBayesianSingleTaskGP
 from botorch.models.fully_bayesian_multitask import SaasFullyBayesianMultiTaskGP
 from botorch.models.gp_regression_fidelity import SingleTaskMultiFidelityGP
 from botorch.models.gp_regression_mixed import MixedSingleTaskGP