From 0d07649d2f55f2794e5d7315b5724d157e83a774 Mon Sep 17 00:00:00 2001
From: saitcakmak <saitcakmak@fb.com>
Date: Mon, 9 Jan 2023 13:37:06 -0800
Subject: [PATCH] Fix TransformedPosterior missing batch shape error in
 _update_base_samples

Summary: See https://github.com/pytorch/botorch/issues/1623. This is the only use case for `posterior.batch_shape`, so fixing it locally makes sense to me.

Differential Revision: D42421494

fbshipit-source-id: 613ab2aa822d53fda615b53979fba857c9b9c931
---
 botorch/sampling/normal.py                    |  8 +++++-
 .../multi_objective/test_monte_carlo.py       | 28 +++++++++++++++++++
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/botorch/sampling/normal.py b/botorch/sampling/normal.py
index 42fb3e4b4c..5223918485 100644
--- a/botorch/sampling/normal.py
+++ b/botorch/sampling/normal.py
@@ -18,6 +18,7 @@
 from botorch.posteriors import Posterior
 from botorch.posteriors.higher_order import HigherOrderGPPosterior
 from botorch.posteriors.multitask import MultitaskGPPosterior
+from botorch.posteriors.transformed import TransformedPosterior
 from botorch.sampling.base import MCSampler
 from botorch.utils.sampling import draw_sobol_normal_samples, manual_seed
 from torch import Tensor
@@ -112,8 +113,13 @@ def _update_base_samples(
                     ..., -n_train_samples:
                 ]
             else:
+                batch_shape = (
+                    posterior._posterior.batch_shape
+                    if isinstance(posterior, TransformedPosterior)
+                    else posterior.batch_shape
+                )
                 single_output = (
-                    len(posterior.base_sample_shape) - len(posterior.batch_shape)
+                    len(posterior.base_sample_shape) - len(batch_shape)
                 ) == 1
                 if single_output:
                     self.base_samples[
diff --git a/test/acquisition/multi_objective/test_monte_carlo.py b/test/acquisition/multi_objective/test_monte_carlo.py
index 99c3eee9fb..b699ea89ae 100644
--- a/test/acquisition/multi_objective/test_monte_carlo.py
+++ b/test/acquisition/multi_objective/test_monte_carlo.py
@@ -37,6 +37,9 @@
 from botorch.models.gp_regression import SingleTaskGP
 from botorch.models.transforms.input import InputPerturbation
 from botorch.models.transforms.outcome import Standardize
+from botorch.posteriors.posterior_list import PosteriorList
+from botorch.posteriors.transformed import TransformedPosterior
+from botorch.sampling.list_sampler import ListSampler
 from botorch.sampling.normal import IIDNormalSampler, SobolQMCNormalSampler
 from botorch.utils.low_rank import sample_cached_cholesky
 from botorch.utils.multi_objective.box_decompositions.dominated import (
@@ -1624,3 +1627,28 @@ def get_acqf(model, matheron):
                     rtol=1e-4,
                 )
             )
+
+    def test_with_transformed(self):
+        # Verify that _set_sampler works with transformed posteriors.
+        mm = MockModel(
+            posterior=PosteriorList(
+                TransformedPosterior(
+                    MockPosterior(samples=torch.rand(2, 3, 1, 1)), lambda X: X
+                ),
+                TransformedPosterior(
+                    MockPosterior(samples=torch.rand(2, 3, 1, 1)), lambda X: X
+                ),
+            )
+        )
+        sampler = ListSampler(
+            IIDNormalSampler(sample_shape=torch.Size([2])),
+            IIDNormalSampler(sample_shape=torch.Size([2])),
+        )
+        # This calls _set_sampler which would previously fail.
+        qNoisyExpectedHypervolumeImprovement(
+            model=mm,
+            ref_point=torch.tensor([0.0, 0.0]),
+            X_baseline=torch.rand(3, 2),
+            sampler=sampler,
+            cache_root=False,
+        )