Fix sample_all_priors to not sample one value for all lengthscales (#2404)

Summary:
Pull Request resolved: #2404

#2371 added support for `sample_all_priors` to handle multivariate priors, but these changes resulted in sampling the same value for all entries of a multi-dimensional hyperparameter if a univariate prior is used. In particular, this means that `sample_all_priors` will sample exactly the same lengthscale for all dimensions when using a univariate prior.

This diff changes this behavior to instead sample according to the shape of the closure when a univariate prior is specified. This results in sampling different lengthscales for each dimension and batch dimension.

Reviewed By: saitcakmak

Differential Revision: D58855726

fbshipit-source-id: 5c70ba83ff1710fe678c60f0ed2f11f57e671ad5

Author: David Eriksson

botorch/optim/utils/model_utils.py

@@ -163,10 +163,16 @@ def sample_all_priors(model: GPyTorchModel, max_retries: int = 100) -> None:
             try:
                 # Set sample shape, so that the prior samples have the same shape
                 # as `closure(module)` without having to be repeated.
-                closure_shape = closure(module).shape
                 prior_shape = prior._extended_shape()
-                sample_shape = closure_shape[: -len(prior_shape)]
-                setting_closure(module, prior.sample(sample_shape=sample_shape))
+                if prior_shape.numel() == 1:
+                    # For a univariate prior we can sample the size of the closure.
+                    # Otherwise we will sample exactly the same value for all
+                    # lengthscales where we commonly specify a univariate prior.
+                    setting_closure(module, prior.sample(closure(module).shape))
+                else:
+                    closure_shape = closure(module).shape
+                    sample_shape = closure_shape[: -len(prior_shape)]
+                    setting_closure(module, prior.sample(sample_shape=sample_shape))
                 break
             except NotImplementedError:
                 warn(

test/optim/utils/test_model_utils.py

@@ -10,6 +10,7 @@
 import warnings
 from copy import deepcopy
 from string import ascii_lowercase
+from typing import Any, Dict
 from unittest.mock import MagicMock, patch
 
 import torch
@@ -246,19 +247,40 @@ def test_sample_all_priors(self):
             with self.assertRaises(RuntimeError):
                 sample_all_priors(model)
 
+    def test_univariate_prior(self) -> None:
+        tkwargs: Dict[str, Any] = {"device": self.device, "dtype": torch.double}
+        for batch in (torch.Size([]), torch.Size([2, 2])):
+            model = SingleTaskGP(
+                train_X=torch.rand(*batch, 5, 3, **tkwargs),
+                train_Y=torch.randn(*batch, 5, 1, **tkwargs),
+                covar_module=RBFKernel(
+                    ard_num_dims=3,
+                    batch_shape=batch,
+                    lengthscale_prior=GammaPrior(6.0, 3.0),  # univariate
+                ),
+            )
+            original_lengthscales = model.covar_module.lengthscale
+            self.assertEqual(len(torch.unique(original_lengthscales)), 1)
+            sample_all_priors(model)
+            new_lengthscales = model.covar_module.lengthscale
+            self.assertFalse(torch.allclose(original_lengthscales, new_lengthscales))
+            # Make sure we sampled different lengthscales (happens with probability 1)
+            self.assertEqual(len(torch.unique(new_lengthscales)), 3 * batch.numel())
+
     def test_with_multivariate_prior(self) -> None:
         # This is modified from https://github.com/pytorch/botorch/issues/780.
+        tkwargs: Dict[str, Any] = {"device": self.device, "dtype": torch.double}
         for batch in (torch.Size([]), torch.Size([3])):
             model = SingleTaskGP(
-                train_X=torch.randn(*batch, 2, 2),
-                train_Y=torch.randn(*batch, 2, 1),
+                train_X=torch.rand(*batch, 2, 2, **tkwargs),
+                train_Y=torch.randn(*batch, 2, 1, **tkwargs),
                 covar_module=RBFKernel(
                     ard_num_dims=2,
                     batch_shape=batch,
                     lengthscale_prior=NormalPrior(
                         # Make this almost singular for easy comparison below.
-                        torch.tensor([[1.0, 1.0]]),
-                        torch.tensor(1e-10),
+                        torch.tensor([[1.0, 1.0]], **tkwargs),
+                        torch.tensor(1e-10, **tkwargs),
                     ),
                 ),
             )
@@ -267,4 +289,4 @@ def test_with_multivariate_prior(self) -> None:
             sample_all_priors(model)
             new_lengthscale = model.covar_module.lengthscale
             self.assertFalse(torch.allclose(original_lengthscale, new_lengthscale))
-            self.assertAllClose(new_lengthscale, torch.ones(*batch, 1, 2))
+            self.assertAllClose(new_lengthscale, torch.ones(*batch, 1, 2, **tkwargs))