Make JAX/numpyro imports lazy in fully Bayesian models (meta-pytorch#3292)

Summary:

`botorch.models.fully_bayesian` and `botorch.models.fully_bayesian_multitask` eagerly import JAX and numpyro at module level. When JAX requires NumPy >= 2.0 but the environment has an older NumPy (e.g., pyper Bento kernel with NumPy 1.24), this causes an `ImportError` even for callers that never use fully Bayesian models.

This diff wraps the JAX/numpyro imports in a `try/except` block, deferring the failure to when `AbstractFullyBayesianSingleTaskGP` or `SaasFullyBayesianMultiTaskGP` is actually instantiated. The error message is preserved and clearly states the NumPy >= 2.0 requirement. Since `from __future__ import annotations` is already present, type annotations referencing `jax.Array` remain valid as strings.

Reviewed By: saitcakmak

Differential Revision: D102367287

Author: sdaulton

botorch/models/fully_bayesian.py

@@ -38,20 +38,17 @@
 from math import log, sqrt
 from typing import Any, TypeVar
 
-import jax.numpy as jnp
 import numpy as np
 
-if np.lib.NumpyVersion(np.__version__) < "2.0.0":
-    raise ImportError(
-        "BoTorch's fully Bayesian models require NumPy >= 2.0 "
-        "(python-scientific-stack version 3). "
-        "Please update your PACKAGE file to set "
-        '"python-scientific-stack": "3".'
-    )
-
-import jax
-import numpyro
-import numpyro.distributions as numpyro_dist
+try:
+    import jax
+    import jax.numpy as jnp
+    import numpyro
+    import numpyro.distributions as numpyro_dist
+
+    _HAS_JAX = True
+except ImportError:  # pragma: no cover
+    _HAS_JAX = False
 import torch
 from botorch.acquisition.objective import PosteriorTransform
 from botorch.models.gpytorch import BatchedMultiOutputGPyTorchModel
@@ -90,6 +87,16 @@
 _sqrt5 = math.sqrt(5)
 
 
+def _check_jax_available() -> None:
+    if not _HAS_JAX:
+        raise ImportError(
+            "BoTorch's fully Bayesian models require JAX, numpyro, and "
+            "NumPy >= 2.0 (python-scientific-stack version 3). "
+            "Please update your PACKAGE file to set "
+            '"python-scientific-stack": "3".'
+        )
+
+
 def matern52_kernel(X: jax.Array, lengthscale: jax.Array) -> jax.Array:
     """Matern-5/2 kernel."""
     dist = compute_dists(X=X, lengthscale=lengthscale)
@@ -872,6 +879,7 @@ def __init__(
             indices_to_warp: An optional list of indices to warp. The default
                 is to warp all inputs.
         """
+        _check_jax_available()
         if not (
             train_X.ndim == train_Y.ndim == 2
             and len(train_X) == len(train_Y)

botorch/models/fully_bayesian_multitask.py

@@ -11,18 +11,22 @@
 from collections.abc import Mapping
 from typing import Any, NoReturn, TypeVar
 
-import jax.numpy as jnp
-import numpyro
-import numpyro.distributions as numpyro_dist
 import torch
 from botorch.acquisition.objective import PosteriorTransform
 from botorch.models.fully_bayesian import (
+    _check_jax_available,
+    _HAS_JAX,
     matern52_kernel,
     MCMC_DIM,
     MIN_INFERRED_NOISE_LEVEL,
     reshape_and_detach,
     SaasPyroModel,
 )
+
+if _HAS_JAX:
+    import jax.numpy as jnp
+    import numpyro
+    import numpyro.distributions as numpyro_dist
 from botorch.models.gpytorch import (
     BatchedMultiOutputGPyTorchModel,
     MultiTaskGPyTorchModel,
@@ -353,6 +357,7 @@ def __init__(
                 input are expected tasks values. If false, unexpected task values
                 will be mapped to the first output_task if supplied.
         """
+        _check_jax_available()
         if not (
             train_X.ndim == train_Y.ndim == 2
             and len(train_X) == len(train_Y)

test/models/test_fully_bayesian.py

@@ -1284,20 +1284,16 @@ class TestFullyBayesianLinearWarpingSingleTaskGP(TestFullyBayesianLinearSingleTa
 
 
 class TestNumpyVersionCheck(BotorchTestCase):
-    def test_old_numpy_raises_import_error(self) -> None:
-        """Test that importing fully_bayesian with NumPy < 2.0 raises ImportError."""
-        import importlib
-        import sys
-
-        from numpy.lib import NumpyVersion
-
-        with patch.object(NumpyVersion, "__lt__", return_value=True):
-            # Remove cached module so the version check re-runs on import.
-            mod_name = "botorch.models.fully_bayesian"
-            saved = sys.modules.pop(mod_name)
-            try:
-                with self.assertRaises(ImportError):
-                    importlib.import_module(mod_name)
-            finally:
-                # Restore the module so other tests are not affected.
-                sys.modules[mod_name] = saved
+    def test_missing_jax_raises_on_instantiation(self) -> None:
+        """Test that missing JAX raises ImportError at model instantiation."""
+        from botorch.models import fully_bayesian
+        from botorch.models.fully_bayesian import _check_jax_available
+
+        with patch.object(fully_bayesian, "_HAS_JAX", False):
+            with self.assertRaises(ImportError):
+                _check_jax_available()
+            with self.assertRaises(ImportError):
+                SaasFullyBayesianSingleTaskGP(
+                    train_X=torch.rand(10, 2),
+                    train_Y=torch.rand(10, 1),
+                )

test/models/test_fully_bayesian_multitask.py

@@ -6,6 +6,7 @@
 
 
 import itertools
+from unittest.mock import patch
 
 import jax.numpy as jnp
 import numpyro.handlers
@@ -182,6 +183,21 @@ def _get_mcmc_samples(self, num_samples: int, dim: int, task_rank: int, **tkwarg
         }
         return mcmc_samples
 
+    def test_missing_jax_raises_on_instantiation(self) -> None:
+        """Test that missing JAX raises ImportError at model instantiation."""
+        from botorch.models import fully_bayesian
+
+        tkwargs = {"device": self.device, "dtype": torch.double}
+        train_X, train_Y, train_Yvar = self._get_base_data(**tkwargs)
+        with patch.object(fully_bayesian, "_HAS_JAX", False):
+            with self.assertRaises(ImportError):
+                SaasFullyBayesianMultiTaskGP(
+                    train_X=train_X,
+                    train_Y=train_Y,
+                    train_Yvar=train_Yvar,
+                    task_feature=4,
+                )
+
     def test_raises(self):
         tkwargs = {"device": self.device, "dtype": torch.double}
         with self.assertRaisesRegex(