Correct handling of input transforms for AdditiveMapSaasSingleTaskGP (#3042)

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

In `AdditiveMapSaasSingleTaskGP`, input transforms that alter the shape of the training data are not handled correctly. This PR fixes it.

A few questions to the Map Saas complex in general:
- We have now the ensemble version and the non-ensemble version, which one would you recommend in general?
- We have getter methods that generate fitted or unfitted versions of the MAP models and we have the option to just initialise the respective class and fit it in the default MAP way. What do you recommend here? Currently, I am using the `AdditiveMapSaasSingleTaskGP` via initialising the class and fitting it like any SingleTaskGP.

### 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: #3042

Test Plan: Unit tests.

Reviewed By: Balandat

Differential Revision: D84059558

Pulled By: saitcakmak

fbshipit-source-id: 415b2ea1ac57c63abb5fa72554cff2cfdd1eae33

Author: jduerholt

botorch/models/map_saas.py

@@ -403,8 +403,15 @@ def __init__(
             if train_Yvar is None
             else None
         )
+        if input_transform is not None:
+            with torch.no_grad():
+                transformed_X = input_transform(train_X)
+            ard_num_dims = transformed_X.shape[-1]
+        else:
+            ard_num_dims = train_X.shape[-1]
+
         covar_module = get_additive_map_saas_covar_module(
-            ard_num_dims=train_X.shape[-1],
+            ard_num_dims=ard_num_dims,
             num_taus=num_taus,
             batch_shape=self._aug_batch_shape,
             # Need to pass dtype and device at initialization of the covar_module

test/models/test_map_saas.py

@@ -7,6 +7,7 @@
 import itertools
 import math
 import pickle
+from functools import partial
 from itertools import product
 from typing import Any
 from unittest import mock
@@ -27,7 +28,12 @@
     get_gaussian_likelihood_with_gamma_prior,
     get_mean_module_with_normal_prior,
 )
-from botorch.models.transforms.input import AppendFeatures, FilterFeatures, Normalize
+from botorch.models.transforms.input import (
+    AppendFeatures,
+    FilterFeatures,
+    Normalize,
+    NumericToCategoricalEncoding,
+)
 from botorch.models.transforms.outcome import Standardize
 from botorch.optim.utils import get_parameters_and_bounds, sample_all_priors
 from botorch.posteriors.fully_bayesian import GaussianMixturePosterior
@@ -43,6 +49,7 @@
 from gpytorch.mlls.exact_marginal_log_likelihood import ExactMarginalLogLikelihood
 from gpytorch.priors import GammaPrior, HalfCauchyPrior, NormalPrior
 from torch import Tensor
+from torch.nn.functional import one_hot
 
 
 class TestMapSaas(BotorchTestCase):
@@ -647,6 +654,30 @@ def _get_data_and_model(
             )
         return train_X, train_Y, train_Yvar, model
 
+    def test_input_transform_dimensions(self) -> None:
+        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)
+
+            input_transform = NumericToCategoricalEncoding(
+                dim=3,
+                categorical_features={2: 3},
+                encoders={2: partial(one_hot, num_classes=3)},
+            )
+
+            model = AdditiveMapSaasSingleTaskGP(
+                train_X=X,
+                train_Y=Y,
+                input_transform=input_transform,
+            )
+            self.assertEqual(model.covar_module.kernels[0].base_kernel.ard_num_dims, 5)
+
     def test_construct_mean_module(self) -> None:
         tkwargs = {"device": self.device, "dtype": torch.double}
         for batch_shape in [None, torch.Size([5])]: