Fix sobol maxdim limitation in prune_baseline (pytorch#419)

Summary:
Fixes facebook/Ax#291

Pull Request resolved: pytorch#419

Test Plan: unit tests

Reviewed By: sdaulton

Differential Revision: D20978955

Pulled By: Balandat

fbshipit-source-id: 2ff30700f1d86542c535f9b39b3c8350e0ec4c25

Author: Balandat

botorch/acquisition/utils.py

@@ -11,12 +11,16 @@
 from __future__ import annotations
 
 import math
+import warnings
 from typing import Callable, Dict, List, Optional
 
 import torch
 from torch import Tensor
+from torch.quasirandom import SobolEngine
 
+from .. import settings
 from ..exceptions.errors import UnsupportedError
+from ..exceptions.warnings import SamplingWarning
 from ..models.model import Model
 from ..sampling.samplers import IIDNormalSampler, SobolQMCNormalSampler
 from ..utils.transforms import squeeze_last_dim
@@ -213,10 +217,19 @@ def prune_inferior_points(
     max_points = math.ceil(max_frac * X.size(-2))
     if max_points < 1 or max_points > X.size(-2):
         raise ValueError(f"max_frac must take values in (0, 1], is {max_frac}")
-    sampler = SobolQMCNormalSampler(num_samples=num_samples)
     with torch.no_grad():
         posterior = model.posterior(X=X)
-        samples = sampler(posterior)
+    if posterior.event_shape.numel() > SobolEngine.MAXDIM:
+        if settings.debug.on():
+            warnings.warn(
+                f"Sample dimension q*m={posterior.event_shape.numel()} exceeding Sobol "
+                f"max dimension ({SobolEngine.MAXDIM}). Using iid samples instead.",
+                SamplingWarning,
+            )
+        sampler = IIDNormalSampler(num_samples=num_samples)
+    else:
+        sampler = SobolQMCNormalSampler(num_samples=num_samples)
+    samples = sampler(posterior)
     if objective is None:
         objective = IdentityMCObjective()
     obj_vals = objective(samples)

test/acquisition/test_utils.py

@@ -5,9 +5,12 @@
 # LICENSE file in the root directory of this source tree.
 
 import itertools
+import warnings
+from contextlib import ExitStack
 from unittest import mock
 
 import torch
+from botorch import settings
 from botorch.acquisition import monte_carlo
 from botorch.acquisition.objective import GenericMCObjective, MCAcquisitionObjective
 from botorch.acquisition.utils import (
@@ -17,7 +20,8 @@
     project_to_target_fidelity,
     prune_inferior_points,
 )
-from botorch.exceptions import UnsupportedError
+from botorch.exceptions.errors import UnsupportedError
+from botorch.exceptions.warnings import SamplingWarning
 from botorch.sampling.samplers import IIDNormalSampler, SobolQMCNormalSampler
 from botorch.utils.testing import BotorchTestCase, MockModel, MockPosterior
 from torch import Tensor
@@ -351,6 +355,22 @@ def test_prune_inferior_points(self):
                 mm = MockModel(MockPosterior(samples=samples))
                 X_pruned = prune_inferior_points(model=mm, X=X)
             self.assertTrue(torch.equal(X_pruned, X[:2]))
+            # test high-dim sampling
+            with ExitStack() as es:
+                mock_event_shape = es.enter_context(
+                    mock.patch(
+                        "botorch.utils.testing.MockPosterior.event_shape",
+                        new_callable=mock.PropertyMock,
+                    )
+                )
+                mock_event_shape.return_value = torch.Size([1, 1, 1112])
+                es.enter_context(
+                    mock.patch.object(MockPosterior, "rsample", return_value=samples)
+                )
+                mm = MockModel(MockPosterior(samples=samples))
+                with warnings.catch_warnings(record=True) as ws, settings.debug(True):
+                    prune_inferior_points(model=mm, X=X)
+                    self.assertTrue(issubclass(ws[-1].category, SamplingWarning))
 
 
 class TestFidelityUtils(BotorchTestCase):