Changes to resampling behavior in MCSamplers (#204)

Summary:
Makes the following changes in case `resample=False`:
1. if `collapse_batch_dims=True` and the requested shape is different from the shape of the base samples, only resample if the last (q and o) dimensions are different, otherwise broadcast the base samples to the requested batch shape (will just be a different number of ones)
2. if `dtype` or `device` of the posterior are different from that of the base samples don't resample if not triggered otherwise - instead automatically move the base samples to the correct device/dtype if apppropriate
Pull Request resolved: #204

Reviewed By: danielrjiang

Differential Revision: D16161405

Pulled By: Balandat

fbshipit-source-id: a7465fd9c4287e0fb9716c23cd10d24e1516e54a

Author: Balandat

botorch/sampling/samplers.py

@@ -87,9 +87,9 @@ def _construct_base_samples(self, posterior: Posterior, shape: torch.Size) -> No
 
          - `resample=True`
          - the MCSampler has no `base_samples` attribute.
-         - `shape` is different than `self.base_samples.shape`.
-         - device and/or dtype of posterior are different than those of
-          `self.base_samples`.
+         - `shape` is different than `self.base_samples.shape` (if
+           `collapse_batch_dims=True`, then batch dimensions of will be
+           automatically broadcasted as necessary)
 
         Args:
             posterior: The Posterior for which to generate base samples.
@@ -137,11 +137,12 @@ def _construct_base_samples(self, posterior: Posterior, shape: torch.Size) -> No
 
         This function will generate a new set of base samples and set the
         `base_samples` buffer if one of the following is true:
-          - `resample=True`
-          - the MCSampler has no `base_samples` attribute.
-          - `shape` is different than `self.base_samples.shape`.
-          - device and/or dtype of posterior ar different than those of
-            `self.base_samples`.
+
+        - `resample=True`
+        - the MCSampler has no `base_samples` attribute.
+        - `shape` is different than `self.base_samples.shape` (if
+            `collapse_batch_dims=True`, then batch dimensions of will be
+            automatically broadcasted as necessary)
 
         Args:
             posterior: The Posterior for which to generate base samples.
@@ -150,16 +151,21 @@ def _construct_base_samples(self, posterior: Posterior, shape: torch.Size) -> No
         if (
             self.resample
             or not hasattr(self, "base_samples")
-            or self.base_samples.shape != shape
-            or self.base_samples.device != posterior.device
-            or self.base_samples.dtype != posterior.dtype
+            or self.base_samples.shape[-2:] != shape[-2:]
+            or (not self.collapse_batch_dims and shape != self.base_samples.shape)
         ):
             with manual_seed(seed=self.seed):
                 base_samples = torch.randn(
                     shape, device=posterior.device, dtype=posterior.dtype
                 )
             self.seed += 1
             self.register_buffer("base_samples", base_samples)
+        elif self.collapse_batch_dims and shape != self.base_samples.shape:
+            self.base_samples = self.base_samples.view(shape)
+        if self.base_samples.device != posterior.device:
+            self.to(device=posterior.device)  # pragma: nocover
+        if self.base_samples.dtype != posterior.dtype:
+            self.to(dtype=posterior.dtype)
 
 
 class SobolQMCNormalSampler(MCSampler):
@@ -201,11 +207,12 @@ def _construct_base_samples(self, posterior: Posterior, shape: torch.Size) -> No
 
         This function will generate a new set of base samples and set the
         `base_samples` buffer if one of the following is true:
-          - `resample=True`
-          - the MCSampler has no `base_samples` attribute.
-          - `self.sample_shape` is different than `self.base_samples.shape`.
-          - device and/or dtype of posterior ar different than those of
-            `self.base_samples`.
+
+        - `resample=True`
+        - the MCSampler has no `base_samples` attribute.
+        - `shape` is different than `self.base_samples.shape` (if
+          `collapse_batch_dims=True`, then batch dimensions of will be
+          automatically broadcasted as necessary)
 
         Args:
             posterior: The Posterior for which to generate base samples.
@@ -214,9 +221,8 @@ def _construct_base_samples(self, posterior: Posterior, shape: torch.Size) -> No
         if (
             self.resample
             or not hasattr(self, "base_samples")
-            or self.base_samples.shape != shape
-            or self.base_samples.device != posterior.device
-            or self.base_samples.dtype != posterior.dtype
+            or self.base_samples.shape[-2:] != shape[-2:]
+            or (not self.collapse_batch_dims and shape != self.base_samples.shape)
         ):
             output_dim = shape[-2:].numel()
             if output_dim > SobolEngine.MAXDIM:
@@ -234,3 +240,9 @@ def _construct_base_samples(self, posterior: Posterior, shape: torch.Size) -> No
             self.seed += 1
             base_samples = base_samples.view(shape)
             self.register_buffer("base_samples", base_samples)
+        elif self.collapse_batch_dims and shape != posterior.event_shape:
+            self.base_samples = self.base_samples.view(shape)
+        if self.base_samples.device != posterior.device:
+            self.to(device=posterior.device)  # pragma: nocover
+        if self.base_samples.dtype != posterior.dtype:
+            self.to(dtype=posterior.dtype)

test/sampling/test_sampler.py

@@ -84,7 +84,13 @@ def test_forward(self, cuda=False):
             posterior_batched = _get_posterior_batched(cuda=cuda, dtype=dtype)
             samples_batched = sampler(posterior_batched)
             self.assertEqual(samples_batched.shape, torch.Size([4, 3, 2, 1]))
-            self.assertEqual(sampler.seed, 1236)
+            self.assertEqual(sampler.seed, 1235)
+            # ensure this works when changing the dtype
+            new_dtype = torch.float if dtype == torch.double else torch.double
+            posterior_batched = _get_posterior_batched(cuda=cuda, dtype=new_dtype)
+            samples_batched = sampler(posterior_batched)
+            self.assertEqual(samples_batched.shape, torch.Size([4, 3, 2, 1]))
+            self.assertEqual(sampler.seed, 1235)
 
             # resample
             sampler = IIDNormalSampler(num_samples=4, resample=True, seed=None)
@@ -212,7 +218,13 @@ def test_forward(self, cuda=False):
             posterior_batched = _get_posterior_batched(cuda=cuda, dtype=dtype)
             samples_batched = sampler(posterior_batched)
             self.assertEqual(samples_batched.shape, torch.Size([4, 3, 2, 1]))
-            self.assertEqual(sampler.seed, 1236)
+            self.assertEqual(sampler.seed, 1235)
+            # ensure this works when changing the dtype
+            new_dtype = torch.float if dtype == torch.double else torch.double
+            posterior_batched = _get_posterior_batched(cuda=cuda, dtype=new_dtype)
+            samples_batched = sampler(posterior_batched)
+            self.assertEqual(samples_batched.shape, torch.Size([4, 3, 2, 1]))
+            self.assertEqual(sampler.seed, 1235)
 
             # resample
             sampler = SobolQMCNormalSampler(num_samples=4, resample=True, seed=None)