Merge branch 'pytorch:main' into fix/fixed_features_dimensionality

Author: AdrianSosic

botorch/models/approximate_gp.py

@@ -399,7 +399,11 @@ def __init__(
                 )
                 train_Y, _ = outcome_transform(train_Y, X=transformed_X)
             self._validate_tensor_args(X=transformed_X, Y=train_Y)
-            validate_input_scaling(train_X=transformed_X, train_Y=train_Y)
+            validate_input_scaling(
+                train_X=transformed_X,
+                train_Y=train_Y,
+                check_nans_only=covar_module is not None,
+            )
             if train_Y.shape[-1] != num_outputs:
                 num_outputs = train_Y.shape[-1]
 

botorch/models/gp_regression.py

@@ -171,6 +171,7 @@ def __init__(
             train_Y=train_Y,
             train_Yvar=train_Yvar,
             ignore_X_dims=ignore_X_dims,
+            check_nans_only=covar_module is not None,
         )
         self._set_dimensions(train_X=train_X, train_Y=train_Y)
         train_X, train_Y, train_Yvar = self._transform_tensor_args(

botorch/models/utils/assorted.py

@@ -228,6 +228,7 @@ def validate_input_scaling(
     train_Yvar: Tensor | None = None,
     raise_on_fail: bool = False,
     ignore_X_dims: list[int] | None = None,
+    check_nans_only: bool = False,
 ) -> None:
     r"""Helper function to validate input data to models.
 
@@ -243,6 +244,10 @@ def validate_input_scaling(
             raised if NaN values are present).
         ignore_X_dims: For this subset of dimensions from `{1, ..., d}`, ignore the
             min-max scaling check.
+        check_nans_only: If True, only check for NaN values. Skips min-max scaling
+            and standardization checks. This is used when the model is provided
+            with a custom covariance module, to avoid potentially irrelevant
+            warnings.
 
     This function is typically called inside the constructor of standard BoTorch
     models. It validates the following:
@@ -261,10 +266,11 @@ def validate_input_scaling(
         check_no_nans(train_Yvar)
         if torch.any(train_Yvar < 0):
             raise InputDataError("Input data contains negative variances.")
-    check_min_max_scaling(
-        X=train_X, raise_on_fail=raise_on_fail, ignore_dims=ignore_X_dims
-    )
-    check_standardization(Y=train_Y, raise_on_fail=raise_on_fail)
+    if not check_nans_only:
+        check_min_max_scaling(
+            X=train_X, raise_on_fail=raise_on_fail, ignore_dims=ignore_X_dims
+        )
+        check_standardization(Y=train_Y, raise_on_fail=raise_on_fail)
 
 
 def mod_batch_shape(module: Module, names: list[str], b: int) -> None:

botorch/optim/optimize.py

@@ -116,7 +116,10 @@ def __post_init__(self) -> None:
                     f"shape is {batch_initial_conditions_shape}."
                 )
 
-            if len(batch_initial_conditions_shape) == 2:
+            if (
+                len(batch_initial_conditions_shape) == 2
+                and self.raw_samples is not None
+            ):
                 warnings.warn(
                     "If using a 2-dim `batch_initial_conditions` botorch will "
                     "default to old behavior of ignoring `num_restarts` and just "
@@ -132,6 +135,7 @@ def __post_init__(self) -> None:
                 len(batch_initial_conditions_shape) == 3
                 and batch_initial_conditions_shape[0] < self.num_restarts
                 and batch_initial_conditions_shape[-2] != self.q
+                and self.raw_samples is not None
             ):
                 warnings.warn(
                     "If using a 3-dim `batch_initial_conditions` where the "

botorch/optim/optimize_mixed.py

@@ -533,6 +533,7 @@ def continuous_step(
         opt_inputs,
         q=1,
         num_restarts=1,
+        raw_samples=None,
         batch_initial_conditions=current_x.unsqueeze(0),
         fixed_features={
             **dict(zip(discrete_dims.tolist(), current_x[discrete_dims])),

test/models/utils/test_assorted.py

@@ -87,7 +87,7 @@ def test_check_min_max_scaling(self):
         X = 0.1 + 0.8 * torch.rand(4, 2, 3)
         with warnings.catch_warnings(record=True) as ws:
             check_min_max_scaling(X=X)
-            self.assertFalse(any(issubclass(w.category, InputDataWarning) for w in ws))
+        self.assertFalse(any(issubclass(w.category, InputDataWarning) for w in ws))
         check_min_max_scaling(X=X, raise_on_fail=True)
         with self.assertWarnsRegex(
             expected_warning=InputDataWarning, expected_regex="not scaled"
@@ -100,30 +100,34 @@ def test_check_min_max_scaling(self):
         Xstd = (X - Xmin) / (Xmax - Xmin)
         with warnings.catch_warnings(record=True) as ws:
             check_min_max_scaling(X=Xstd)
-            self.assertFalse(any(issubclass(w.category, InputDataWarning) for w in ws))
+        self.assertFalse(any(issubclass(w.category, InputDataWarning) for w in ws))
         check_min_max_scaling(X=Xstd, raise_on_fail=True)
         with warnings.catch_warnings(record=True) as ws:
             check_min_max_scaling(X=Xstd, strict=True)
-            self.assertFalse(any(issubclass(w.category, InputDataWarning) for w in ws))
+        self.assertFalse(any(issubclass(w.category, InputDataWarning) for w in ws))
         check_min_max_scaling(X=Xstd, strict=True, raise_on_fail=True)
         # check violation
         X[0, 0, 0] = 2
         with warnings.catch_warnings(record=True) as ws:
             check_min_max_scaling(X=X)
-            self.assertTrue(any(issubclass(w.category, InputDataWarning) for w in ws))
-            self.assertTrue(any("not contained" in str(w.message) for w in ws))
+        self.assertTrue(any(issubclass(w.category, InputDataWarning) for w in ws))
+        self.assertTrue(any("not contained" in str(w.message) for w in ws))
         with self.assertRaises(InputDataError):
             check_min_max_scaling(X=X, raise_on_fail=True)
         with warnings.catch_warnings(record=True) as ws:
             check_min_max_scaling(X=X, strict=True)
-            self.assertTrue(any(issubclass(w.category, InputDataWarning) for w in ws))
-            self.assertTrue(any("not contained" in str(w.message) for w in ws))
+        self.assertTrue(any(issubclass(w.category, InputDataWarning) for w in ws))
+        self.assertTrue(any("not contained" in str(w.message) for w in ws))
         with self.assertRaises(InputDataError):
             check_min_max_scaling(X=X, strict=True, raise_on_fail=True)
         # check ignore_dims
         with warnings.catch_warnings(record=True) as ws:
             check_min_max_scaling(X=X, ignore_dims=[0])
-            self.assertFalse(any(issubclass(w.category, InputDataWarning) for w in ws))
+        self.assertFalse(any(issubclass(w.category, InputDataWarning) for w in ws))
+        # all dims ignored
+        with warnings.catch_warnings(record=True) as ws:
+            check_min_max_scaling(X=X, ignore_dims=[0, 1, 2])
+        self.assertFalse(any(issubclass(w.category, InputDataWarning) for w in ws))
 
     def test_check_standardization(self):
         # Ensure that it is not filtered out.
@@ -181,6 +185,11 @@ def test_validate_input_scaling(self):
         # check that errors are raised when requested
         with self.assertRaises(InputDataError):
             validate_input_scaling(train_X=train_X, train_Y=train_Y, raise_on_fail=True)
+        # check that normalization & standardization checks & errors are skipped when
+        # check_nans_only is True
+        validate_input_scaling(
+            train_X=train_X, train_Y=train_Y, raise_on_fail=True, check_nans_only=True
+        )
         # check that no errors are being raised if everything is standardized
         train_X_min = train_X.min(dim=-1, keepdim=True)[0]
         train_X_max = train_X.max(dim=-1, keepdim=True)[0]
@@ -202,6 +211,11 @@ def test_validate_input_scaling(self):
         train_X_std[0, 0, 0] = float("nan")
         with self.assertRaises(InputDataError):
             validate_input_scaling(train_X=train_X_std, train_Y=train_Y_std)
+        # NaNs still raise errors when check_nans_only is True
+        with self.assertRaises(InputDataError):
+            validate_input_scaling(
+                train_X=train_X_std, train_Y=train_Y_std, check_nans_only=True
+            )
 
 
 class TestGPTPosteriorSettings(BotorchTestCase):