Make fully Bayesian imports lazy (meta-pytorch#3265)

Summary:
X-link: facebook/Ax#5148

Pull Request resolved: meta-pytorch#3265

`import botorch` eagerly loaded Pyro via two import paths:
1. `botorch.fit` → `from pyro.infer.mcmc import MCMC, NUTS`
2. `botorch.models.__init__` → `fully_bayesian.py` → Pyro imports

This diff makes fully Bayesian imports lazy so that `import botorch` no longer
eagerly loads Pyro. The imports are deferred until `fit_fully_bayesian_model_nuts`
is called or `SaasFullyBayesianSingleTaskGP`/`SaasFullyBayesianMultiTaskGP` are
accessed from `botorch.models`.

Reviewed By: saitcakmak

Differential Revision: D99687047

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,
@@ -49,6 +49,14 @@
 gp_settings.max_eager_kernel_size._global_value = 4096
 
 
+def __getattr__(name: str):
+    if name == "fit_fully_bayesian_model_nuts":
+        from botorch.fit import fit_fully_bayesian_model_nuts
+
+        return fit_fully_bayesian_model_nuts
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
 __all__ = [
     "acquisition",
     "batch_cross_validation",

botorch/fit.py

@@ -12,15 +12,13 @@
 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
 
 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
@@ -43,11 +41,14 @@
 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 pyro.infer.mcmc import MCMC, NUTS
 from torch import device, Tensor
 from torch.nn import Parameter
 from torch.utils.data import DataLoader
 
+if TYPE_CHECKING:
+    from botorch.models.fully_bayesian import AbstractFullyBayesianSingleTaskGP
+    from botorch.models.fully_bayesian_multitask import SaasFullyBayesianMultiTaskGP
+
 
 def _debug_warn(w: WarningMessage) -> bool:
     if _LBFGSB_MAXITER_MAXFUN_REGEX.search(str(w.message)):
@@ -363,6 +364,8 @@ def fit_fully_bayesian_model_nuts(
         >>> gp = SaasFullyBayesianSingleTaskGP(train_X, train_Y)
         >>> fit_fully_bayesian_model_nuts(gp)
     """
+    from pyro.infer.mcmc 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
@@ -30,6 +28,21 @@
 from botorch.models.multitask import KroneckerMultiTaskGP, MultiTaskGP
 from botorch.models.pairwise_gp import PairwiseGP, PairwiseLaplaceMarginalLogLikelihood
 
+_FULLY_BAYESIAN_LAZY_IMPORTS = {
+    "SaasFullyBayesianSingleTaskGP": "botorch.models.fully_bayesian",
+    "SaasFullyBayesianMultiTaskGP": "botorch.models.fully_bayesian_multitask",
+}
+
+
+def __getattr__(name: str):
+    if name in _FULLY_BAYESIAN_LAZY_IMPORTS:
+        import importlib
+
+        module = importlib.import_module(_FULLY_BAYESIAN_LAZY_IMPORTS[name])
+        return getattr(module, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
 __all__ = [
     "add_saas_prior",
     "AdditiveMapSaasSingleTaskGP",