Fix bug in condition_on_observation (#3034)

Summary:
<!--
Thank you for sending the PR! We appreciate you spending the time to make BoTorch better.

Help us understand your motivation by explaining why you decided to make this change.

You can learn more about contributing to BoTorch here: https://github.com/meta-pytorch/botorch/blob/main/CONTRIBUTING.md
-->

## Motivation

Fix the bug outlined here: #3033

Tests are not yet adapted,  first, I want to get your opinion on the fix ;)

### Have you read the [Contributing Guidelines on pull requests](https://github.com/meta-pytorch/botorch/blob/main/CONTRIBUTING.md#pull-requests)?

Yes.

Pull Request resolved: #3034

Test Plan: Unit tests, not yet adapted.

Reviewed By: hvarfner

Differential Revision: D83766789

Pulled By: saitcakmak

fbshipit-source-id: 787d885d617e398fae22c2447ba76156b75f5d46

Author: jduerholt

botorch/models/gpytorch.py

@@ -244,6 +244,7 @@ def condition_on_observations(
         """
         # pass the transformed data to get_fantasy_model below
         # (unless we've already transformed if BatchedMultiOutputGPyTorchModel)
+        X_original = X.clone()
         X = self.transform_inputs(X)
 
         Yvar = noise
@@ -270,9 +271,9 @@ def condition_on_observations(
         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_original.shape[:-2], fantasy_model._original_train_inputs.shape[:-2]
             )
-            X_expanded = X.expand(expand_shape + X.shape[-2:])
+            X_expanded = X_original.expand(expand_shape + X_original.shape[-2:])
             orig_expanded = fantasy_model._original_train_inputs.expand(
                 expand_shape + fantasy_model._original_train_inputs.shape[-2:]
             )

test/models/test_gpytorch.py

@@ -6,6 +6,7 @@
 
 import itertools
 import warnings
+from functools import partial
 
 import torch
 from botorch.acquisition.objective import ScalarizedPosteriorTransform
@@ -24,7 +25,11 @@
 from botorch.models.model import FantasizeMixin
 from botorch.models.multitask import MultiTaskGP
 from botorch.models.transforms import Standardize
-from botorch.models.transforms.input import ChainedInputTransform, InputTransform
+from botorch.models.transforms.input import (
+    ChainedInputTransform,
+    InputTransform,
+    NumericToCategoricalEncoding,
+)
 from botorch.models.utils import fantasize
 from botorch.posteriors.gpytorch import GPyTorchPosterior
 from botorch.sampling.normal import SobolQMCNormalSampler
@@ -39,6 +44,8 @@
 from gpytorch.settings import trace_mode
 from torch import Tensor
 
+from torch.nn.functional import one_hot
+
 
 class SimpleInputTransform(InputTransform, torch.nn.Module):
     def __init__(self, transform_on_train: bool) -> None:
@@ -691,6 +698,43 @@ def test_condition_on_observations_model_list(self):
                     X=torch.rand(2, 1, **tkwargs), Y=torch.rand(2, 2, **tkwargs)
                 )
 
+    def test_condition_on_observations_input_transform_shape_manipulation(self):
+        for dtype in (torch.float, torch.double):
+            tkwargs = {"device": self.device, "dtype": dtype}
+
+            # Create data
+            X = torch.rand(12, 2, **tkwargs) * 2
+            Y = 1 - (X - 0.5).norm(dim=-1, keepdim=True)
+            Y += 0.1 * torch.rand_like(Y)
+            # Add a categorical feature
+            new_col = torch.randint(0, 3, (X.shape[0], 1), **tkwargs)
+            X = torch.cat([X, new_col], dim=1)
+
+            train_X = X[:10]
+            train_Y = Y[:10]
+
+            condition_X = X[10:]
+            condition_Y = Y[10:]
+
+            # setup transform and model
+            input_transform = NumericToCategoricalEncoding(
+                dim=3,
+                categorical_features={2: 3},
+                encoders={2: partial(one_hot, num_classes=3)},
+            )
+
+            model = SimpleGPyTorchModel(
+                train_X, train_Y, input_transform=input_transform
+            )
+            model.eval()
+            _ = model.posterior(train_X)
+
+            conditioned_model = model.condition_on_observations(
+                condition_X, condition_Y
+            )
+            self.assertAllClose(conditioned_model._original_train_inputs, X)
+            self.assertAllClose(conditioned_model.train_inputs[0], input_transform(X))
+
     def test_condition_on_observations_input_transform_consistency(self):
         """Test that input transforms are applied consistently in
          condition_on_observations.