be explicit about task feature in TL (#2918)

Summary:
Pull Request resolved: #2918

X-link: facebook/Ax#4021

Previously, the TL adapter added a task feature to `SearchSpaceDigest.task_features`, but not to `SearchSpaceDigest.bounds` or `SearchSpaceDigest.feature_names`. This avoided needing to pass the task feature as a fixed_feature when making predictions or generating.

However, there was a bug when using `Normalize` that dropped one non-task feature here: https://github.com/facebook/Ax/blob/main/ax/generators/torch/botorch_modular/surrogate.py#L247-L250

This diff adds the task feature to `SearchSpaceDigest.bounds`,  `SearchSpaceDigest.feature_names`, and `MultiTaskDataset.feature_names` in order to be explicity about all features/bounds. This requires also adding the task feature as a fixed feature before generating, predicting, and cross-validating.

Reviewed By: saitcakmak

Differential Revision: D77905821

Author: sdaulton

botorch/utils/datasets.py

@@ -18,6 +18,8 @@
 from pyre_extensions import none_throws
 from torch import long, ones, Tensor
 
+TASK_FEATURE_NAME = "task_feature"
+
 
 class SupervisedDataset:
     r"""Base class for datasets consisting of labelled pairs `(X, Y)`
@@ -351,7 +353,9 @@ def __init__(
         self.target_outcome_name = target_outcome_name
         self.task_feature_index = task_feature_index
         self._validate_datasets(datasets=datasets)
-        self.feature_names = self.datasets[target_outcome_name].feature_names
+        self.feature_names = self.datasets[target_outcome_name].feature_names.copy()
+        if task_feature_index is None:
+            self.feature_names.append(TASK_FEATURE_NAME)
         self.outcome_names = [target_outcome_name]
 
         # Check if the datasets have identical feature sets.