Wire LearnedFeatureImputation and map_heterogeneous_to_full for MultiTaskGP

Summary:
X-link: facebook/Ax#5192

Automatically configures learned feature imputation for models that pad heterogeneous per-task data to the full joint feature space. Models with native heterogeneity support are excluded from this automatic configuration.

Differential Revision: D101841497

Author: Carl Hvarfner

botorch/models/heterogeneous_mtgp.py

@@ -312,6 +312,7 @@ def construct_inputs(
         rank: int | None = None,
         use_saas_prior: bool = True,
         use_combinatorial_kernel: bool = True,
+        map_heterogeneous_to_full: bool = False,
     ) -> dict[str, Any]:
         r"""Construct ``Model`` keyword arguments from a given ``MultiTaskDataset``.
 
@@ -326,6 +327,10 @@ def construct_inputs(
                 ``MultiTaskConditionalKernel``.
             use_combinatorial_kernel: Whether to use a combinatorial kernel over the
                 binary embedding of task features in ``MultiTaskConditionalKernel``.
+            map_heterogeneous_to_full: Accepted for compatibility with
+                ``MultiTaskGP.construct_inputs`` but unused.
+                ``HeterogeneousMTGP`` handles heterogeneous features via
+                ``MultiTaskConditionalKernel``.
         """
         if training_data.task_feature_index != -1:
             raise NotImplementedError(

botorch/models/transforms/input.py

@@ -2051,12 +2051,16 @@ def __init__(
             ),
         )
         if bounds is not None:
-            # Pad bounds with dummy [0, 1] for the task column so the Interval
-            # constraint has shape (d+1,) matching raw_imputation_values.
-            padded_lower = torch.zeros(d + 1, dtype=dtype, device=device)
-            padded_upper = torch.ones(d + 1, dtype=dtype, device=device)
-            padded_lower[:d] = bounds[0]
-            padded_upper[:d] = bounds[1]
+            # Constraint bounds must match raw_imputation_values' (num_tasks,
+            # d+1) shape — gpytorch's scipy fitting path flattens the parameter
+            # and expects per-element bounds, so a (d+1,) bound would fail to
+            # broadcast. Pad the task column slot with dummy [0, 1].
+            padded_lower = torch.zeros(
+                self.num_tasks, d + 1, dtype=dtype, device=device
+            )
+            padded_upper = torch.ones(self.num_tasks, d + 1, dtype=dtype, device=device)
+            padded_lower[:, :d] = bounds[0]
+            padded_upper[:, :d] = bounds[1]
             self.register_constraint(
                 "raw_imputation_values",
                 Interval(

test/models/test_heterogeneous_mtgp.py

@@ -97,6 +97,12 @@ def test_input_constructor(self) -> None:
         self.assertEqual(model_inputs["full_feature_dim"], 5)
         self.assertIsNone(model_inputs["rank"])
 
+        with self.subTest("map_heterogeneous_to_full accepted and ignored"):
+            model_inputs = HeterogeneousMTGP.construct_inputs(
+                training_data=self.mtds, map_heterogeneous_to_full=True
+            )
+            self.assertNotIn("map_heterogeneous_to_full", model_inputs)
+
     def test_standard_heterogeneous_mtgp(self) -> None:
         # Construct the model (inferred noise: train_Yvars is None).
         model_inputs = HeterogeneousMTGP.construct_inputs(training_data=self.mtds)