Reap base_samples argument of GPyTorchPosterior.rsample (#2254)

Summary:
X-link: facebookexternal/botorch_fb#17


The `base_samples` argument to GPyTorchPosterior.rsample was deprecated in BoTorch 8.0.0
- reaped the deprecated code
- removed the corresponding unit tests
- Removed `base_samples` argument from `MockPosterior.rsample` used for testing, since the base Posterior class also does not permit `base_samples` in `rsample`. Moved an exception in `MockPosterior` so that invalid `base_samples` will still be checked.
- fixed a docstring in PosteriorList

Reviewed By: saitcakmak

Differential Revision: D55038429

Author: esantorella

botorch/posteriors/gpytorch.py

@@ -146,11 +146,7 @@ def rsample_from_base_samples(
             samples = samples.unsqueeze(-1)
         return samples
 
-    def rsample(
-        self,
-        sample_shape: Optional[torch.Size] = None,
-        base_samples: Optional[Tensor] = None,
-    ) -> Tensor:
+    def rsample(self, sample_shape: Optional[torch.Size] = None) -> Tensor:
         r"""Sample from the posterior (with gradients).
 
         Args:
@@ -167,30 +163,12 @@ def rsample(
         """
         if sample_shape is None:
             sample_shape = torch.Size([1])
-        if base_samples is not None:
-            warnings.warn(
-                "Use of `base_samples` with `rsample` is deprecated. Use "
-                "`rsample_from_base_samples` instead.",
-                DeprecationWarning,
-            )
-            if base_samples.shape[: len(sample_shape)] != sample_shape:
-                raise RuntimeError(
-                    "`sample_shape` disagrees with shape of `base_samples`. "
-                    f"Got {sample_shape=} and {base_samples.shape=}."
-                )
-            # get base_samples to the correct shape
-            base_samples = base_samples.expand(self._extended_shape(sample_shape))
-            if not self._is_mt:
-                # Remove output dimension in single output case.
-                base_samples = base_samples.squeeze(-1)
-            return self.rsample_from_base_samples(
-                sample_shape=sample_shape, base_samples=base_samples
-            )
+
         with ExitStack() as es:
             if linop_settings._fast_covar_root_decomposition.is_default():
                 es.enter_context(linop_settings._fast_covar_root_decomposition(False))
             samples = self.distribution.rsample(
-                sample_shape=sample_shape, base_samples=base_samples
+                sample_shape=sample_shape, base_samples=None
             )
         # make sure there always is an output dimension
         if not self._is_mt:

botorch/posteriors/posterior_list.py

@@ -154,19 +154,13 @@ def variance(self) -> Tensor:
         """
         return self._reshape_and_cat(tensors=[p.variance for p in self.posteriors])
 
-    def rsample(
-        self,
-        sample_shape: Optional[torch.Size] = None,
-    ) -> Tensor:
+    def rsample(self, sample_shape: Optional[torch.Size] = None) -> Tensor:
         r"""Sample from the posterior (with gradients).
 
         Args:
             sample_shape: A `torch.Size` object specifying the sample shape. To
                 draw `n` samples, set to `torch.Size([n])`. To draw `b` batches
                 of `n` samples each, set to `torch.Size([b, n])`.
-            base_samples: An (optional) Tensor of `N(0, I)` base samples of
-                appropriate dimension, typically obtained from a `Sampler`.
-                This is used for deterministic optimization. Deprecated.
 
         Returns:
             Samples from the posterior, a tensor of shape

botorch/utils/testing.py

@@ -298,24 +298,24 @@ def variance(self):
     def rsample(
         self,
         sample_shape: Optional[torch.Size] = None,
-        base_samples: Optional[Tensor] = None,
     ) -> Tensor:
         """Mock sample by repeating self._samples. If base_samples is provided,
         do a shape check but return the same mock samples."""
         if sample_shape is None:
             sample_shape = torch.Size()
-        if sample_shape is not None and base_samples is not None:
-            # check the base_samples shape is consistent with the sample_shape
-            if base_samples.shape[: len(sample_shape)] != sample_shape:
-                raise RuntimeError("sample_shape disagrees with base_samples.")
         return self._samples.expand(sample_shape + self._samples.shape)
 
     def rsample_from_base_samples(
         self,
         sample_shape: torch.Size,
         base_samples: Tensor,
     ) -> Tensor:
-        return self.rsample(sample_shape, base_samples)
+        if base_samples.shape[: len(sample_shape)] != sample_shape:
+            raise RuntimeError(
+                "`sample_shape` disagrees with shape of `base_samples`. "
+                f"Got {sample_shape=} and {base_samples.shape=}."
+            )
+        return self.rsample(sample_shape)
 
 
 @GetSampler.register(MockPosterior)

test/posteriors/test_gpytorch.py

@@ -80,10 +80,6 @@ def test_GPyTorchPosterior(self):
                 posterior.rsample_from_base_samples(
                     sample_shape=torch.Size([3]), base_samples=base_samples
                 )
-            with self.assertRaisesRegex(RuntimeError, "sample_shape"):
-                posterior.rsample(
-                    sample_shape=torch.Size([3]), base_samples=base_samples
-                )
             # ensure consistent result
             for sample_shape in ([4], [4, 2]):
                 base_samples = torch.randn(
@@ -109,20 +105,6 @@ def test_GPyTorchPosterior(self):
             self.assertEqual(density.shape, posterior._extended_shape(torch.Size([2])))
             expected = torch.stack([marginal.log_prob(q).exp() for q in q_val], dim=0)
             self.assertAllClose(density, expected)
-            # collapse_batch_dims
-            b_mean = torch.rand(2, 3, dtype=dtype, device=self.device)
-            b_variance = 1 + torch.rand(2, 3, dtype=dtype, device=self.device)
-            b_covar = torch.diag_embed(b_variance)
-            b_mvn = MultivariateNormal(b_mean, to_linear_operator(b_covar))
-            b_posterior = GPyTorchPosterior(distribution=b_mvn)
-            b_base_samples = torch.randn(4, 1, 3, 1, device=self.device, dtype=dtype)
-            with self.assertWarnsRegex(
-                DeprecationWarning, "`base_samples` with `rsample`"
-            ):
-                b_samples = b_posterior.rsample(
-                    sample_shape=torch.Size([4]), base_samples=b_base_samples
-                )
-            self.assertEqual(b_samples.shape, torch.Size([4, 2, 3, 1]))
 
     def test_GPyTorchPosterior_Multitask(self):
         for dtype in (torch.float, torch.double):
@@ -159,20 +141,6 @@ def test_GPyTorchPosterior_Multitask(self):
                 sample_shape=torch.Size([4, 2]), base_samples=base_samples2
             )
             self.assertAllClose(samples2_b1, samples2_b2)
-            # collapse_batch_dims
-            b_mean = torch.rand(2, 3, 2, dtype=dtype, device=self.device)
-            b_variance = 1 + torch.rand(2, 3, 2, dtype=dtype, device=self.device)
-            b_covar = torch.diag_embed(b_variance.view(2, 6))
-            b_mvn = MultitaskMultivariateNormal(b_mean, to_linear_operator(b_covar))
-            b_posterior = GPyTorchPosterior(distribution=b_mvn)
-            b_base_samples = torch.randn(4, 1, 3, 2, device=self.device, dtype=dtype)
-            with self.assertWarnsRegex(
-                DeprecationWarning, "`base_samples` with `rsample`"
-            ):
-                b_samples = b_posterior.rsample(
-                    sample_shape=torch.Size([4]), base_samples=b_base_samples
-                )
-            self.assertEqual(b_samples.shape, torch.Size([4, 2, 3, 2]))
 
     def test_degenerate_GPyTorchPosterior(self):
         for dtype, multi_task in (

test/utils/test_testing.py

@@ -8,17 +8,17 @@
 from botorch.utils.testing import BotorchTestCase, MockModel, MockPosterior
 
 
-class TestMock(BotorchTestCase):
-    def test_MockPosterior(self):
-        # test basic logic
+class TestMockPosterior(BotorchTestCase):
+    def test_basic_logic(self) -> None:
         mp = MockPosterior()
         self.assertEqual(mp.device.type, "cpu")
         self.assertEqual(mp.dtype, torch.float32)
         self.assertEqual(mp._extended_shape(), torch.Size())
         self.assertEqual(
             MockPosterior(variance=torch.rand(2))._extended_shape(), torch.Size([2])
         )
-        # test passing in tensors
+
+    def test_passing_tensors(self) -> None:
         mean = torch.rand(2)
         variance = torch.eye(2)
         samples = torch.rand(1, 2)
@@ -31,10 +31,17 @@ def test_MockPosterior(self):
         self.assertTrue(
             torch.all(mp.rsample(torch.Size([2])) == samples.repeat(2, 1, 1))
         )
-        with self.assertRaises(RuntimeError):
-            mp.rsample(sample_shape=torch.Size([2]), base_samples=torch.rand(3))
 
-    def test_MockModel(self):
+    def test_rsample_from_base_samples(self) -> None:
+        mp = MockPosterior()
+        with self.assertRaisesRegex(
+            RuntimeError, "`sample_shape` disagrees with shape of `base_samples`."
+        ):
+            mp.rsample_from_base_samples(torch.zeros(2, 2), torch.zeros(3))
+
+
+class TestMockModel(BotorchTestCase):
+    def test_basic(self) -> None:
         mp = MockPosterior()
         mm = MockModel(mp)
         X = torch.empty(0)