Condition_on_observations adds data to train_inputs when input_transforms are applied (#2990)

Summary:
Pull Request resolved: #2990

When applying input transforms (even with transform_on_train=True), conditioning on new data does not add it to train_inputs. This is unintuitive, as it is not clear why input transforms should disallow the user from adding observations to the training data this way.

Since condition_on_observations returns a model (and thus, the user has the flexibility to override the old model or not with the fantasy model), the data should always be added to training.

Fixes #2533.

Simple notebook on why this doesn't make sense:
N7899660

Reviewed By: saitcakmak

Differential Revision: D80830741

fbshipit-source-id: 422abfa9d9d6212e52e15fa25e445a73d0b3dc42

Author: Carl Hvarfner

botorch/models/gpytorch.py

@@ -242,11 +242,13 @@ def condition_on_observations(
             >>> new_Y = torch.sin(new_X[:, :1]) + torch.cos(new_X[:, 1:])
             >>> model = model.condition_on_observations(X=new_X, Y=new_Y)
         """
+        # pass the transformed data to get_fantasy_model below
+        # (unless we've already transformed if BatchedMultiOutputGPyTorchModel)
         X = self.transform_inputs(X)
+
         Yvar = noise
         if hasattr(self, "outcome_transform"):
-            # pass the transformed data to get_fantasy_model below
-            # (unless we've already trasnformed if BatchedMultiOutputGPyTorchModel)
+            # And do the same for the outcome transform, if it exists.
             if not isinstance(self, BatchedMultiOutputGPyTorchModel):
                 # `noise` is assumed to already be outcome-transformed.
                 Y, _ = self.outcome_transform(Y=Y, Yvar=Yvar, X=X)
@@ -260,9 +262,25 @@ def condition_on_observations(
             if Yvar is not None:
                 kwargs.update({"noise": Yvar.squeeze(-1)})
         # get_fantasy_model will properly copy any existing outcome transforms
-        # (since it deepcopies the original model)
-
-        return self.get_fantasy_model(inputs=X, targets=Y, **kwargs)
+        # (since it deepcopies the original model))
+        fantasy_model = self.get_fantasy_model(inputs=X, targets=Y, **kwargs)
+
+        # If we use an input transform, the fantasized data will not get added to
+        # the training data by default. We need to manually add it.
+        if hasattr(fantasy_model, "input_transform"):
+            # Broadcast tensors to compatible shape before concatenating
+            expand_shape = torch.broadcast_shapes(
+                X.shape[:-2], fantasy_model._original_train_inputs.shape[:-2]
+            )
+            X_expanded = X.expand(expand_shape + X.shape[-2:])
+            orig_expanded = fantasy_model._original_train_inputs.expand(
+                expand_shape + fantasy_model._original_train_inputs.shape[-2:]
+            )
+            fantasy_model._original_train_inputs = torch.cat(
+                [orig_expanded, X_expanded],
+                dim=-2,
+            ).detach()
+        return fantasy_model
 
 
 # pyre-fixme[13]: uninitialized attributes _num_outputs, _input_batch_shape,

test/models/test_fully_bayesian.py

@@ -792,7 +792,7 @@ def test_construct_inputs(self) -> None:
             else:
                 self.assertTrue(Yvar.equal(data_dict["train_Yvar"]))
 
-    def test_fbstgp_condition_on_observations(self) -> None:
+    def test_condition_on_observations(self) -> None:
         # The following conditioned data shapes should work (output describes):
         # training data shape after cond(batch shape in output is req. in gpytorch)
         # X: num_models x n x d, Y: num_models x n x d --> num_models x n x d

test/models/test_gpytorch.py

@@ -734,3 +734,95 @@ def test_condition_on_observations_input_transform_consistency(self):
                 conditioned_model.train_inputs[0],
                 expected_combined_inputs,
             )
+
+    def test_condition_on_observations_train_input_shapes(self):
+        """Comprehensive test for condition_on_observations functionality.
+
+        Tests input transform consistency, train/eval mode stability,
+        different transform settings, and batch shape handling.
+        """
+        for dtype in (torch.float, torch.double):
+            tkwargs = {"device": self.device, "dtype": dtype}
+
+            # Test 1: Train/eval mode stability with transform_on_train=True
+            train_X = torch.tensor([[0.0], [1.0]], **tkwargs)
+            train_Y = torch.tensor([[1.0], [2.0]], **tkwargs)
+            input_transform = SimpleInputTransform(transform_on_train=True)
+            model = SimpleGPyTorchModel(
+                train_X, train_Y, input_transform=input_transform
+            )
+
+            new_X = torch.tensor([[0.5]], **tkwargs)
+            new_Y = torch.tensor([[1.5]], **tkwargs)
+            _ = model.posterior(train_X)
+            conditioned_model = model.condition_on_observations(new_X, new_Y)
+
+            # Verify conditioned observations persist across train/eval modes
+            conditioned_model.eval()
+            self.assertEqual(conditioned_model.train_targets.shape[0], 3)
+            conditioned_model.train()
+            self.assertEqual(conditioned_model.train_targets.shape[0], 3)
+            self.assertEqual(conditioned_model.train_inputs[0].shape[0], 3)
+
+            # Test 2: Transform behavior with transform_on_train=False
+            model2 = SimpleGPyTorchModel(
+                train_X,
+                train_Y,
+                input_transform=SimpleInputTransform(transform_on_train=False),
+            )
+            _ = model2.posterior(train_X)
+            conditioned_model2 = model2.condition_on_observations(new_X, new_Y)
+            self.assertEqual(conditioned_model2.train_targets.shape[0], 3)
+
+            # Verify model can make predictions after conditioning
+            test_X = torch.tensor([[0.25]], **tkwargs)
+            posterior = conditioned_model2.posterior(test_X)
+            self.assertEqual(posterior.mean.shape, torch.Size([1, 1]))
+
+        # Test 3: Batch shape handling and broadcasting (double precision only)
+        tkwargs = {"device": self.device, "dtype": torch.double}
+
+        # Same ndim - should update _original_train_inputs
+        train_X = torch.rand(2, 2, **tkwargs)
+        train_Y = torch.rand(2, 1, **tkwargs)
+        model = SimpleGPyTorchModel(
+            train_X, train_Y, input_transform=SimpleInputTransform(True)
+        )
+        _ = model.posterior(train_X)
+
+        original_size = model._original_train_inputs.shape[0]
+        fantasy_model = model.condition_on_observations(
+            torch.rand(1, 2, **tkwargs), torch.rand(1, 1, **tkwargs)
+        )
+        self.assertEqual(
+            fantasy_model._original_train_inputs.shape[0], original_size + 1
+        )
+
+        # Different ndim - should NOT update _original_train_inputs
+        original_size = model._original_train_inputs.shape[0]
+        fantasy_model = model.condition_on_observations(
+            torch.rand(3, 2, 2, **tkwargs), torch.rand(3, 2, 1, **tkwargs)
+        )
+
+        # NOTE expected behavior is expand (2, 2) & (3, 2, 2) is expanded
+        # and then concatenated along dimension -2 --> (3, 4, 2)
+        self.assertEqual(
+            fantasy_model._original_train_inputs.shape, torch.Size([3, 4, 2])
+        )
+
+        # Test 4: Fantasy model behavior
+        model2 = SimpleGPyTorchModel(
+            train_X, train_Y, input_transform=SimpleInputTransform(True)
+        )
+        _ = model2.posterior(train_X)
+        original_size = model2._original_train_inputs.shape[0]
+
+        fantasy_model = model2.condition_on_observations(
+            torch.rand(1, 2, **tkwargs), torch.rand(1, 1, **tkwargs)
+        )
+
+        # Fantasy model gets data, original model does not
+        self.assertEqual(
+            fantasy_model._original_train_inputs.shape[0], original_size + 1
+        )
+        self.assertEqual(model2._original_train_inputs.shape[0], original_size)