Simplify task value remapping API (#4860)

Summary:
X-link: meta-pytorch/botorch#3163

Pull Request resolved: #4860

X-link: meta-pytorch/botorch#3146

Simplifies the get_task_value_remapping() API from 4 parameters to 2, addressing confusion reported in #3085.

The observed_task_values parameter is removed because the parent diff (D90769576) now makes MultiTaskGP track observed/unobserved tasks internally via _observed_task_indices and _unobserved_task_indices. The default_task_value parameter is removed because the previous behavior—silently mapping unknown tasks to an arbitrary fallback—was confusing and error-prone; instead, unrecognized tasks now map to NaN, providing an explicit error sentinel with a clear warning message.

Differential Revision: D90998243

Author: Carl Hvarfner

ax/generators/torch/botorch_modular/input_constructors/outcome_transform.py

@@ -106,17 +106,10 @@ def _outcome_transform_argparse_stratified_standardize(
     dataset = assert_is_instance(dataset, MultiTaskDataset)
     if dataset.has_heterogeneous_features:
         task_feature_index = dataset.task_feature_index or -1
-        task_values = torch.arange(len(dataset.datasets), dtype=torch.long)
     else:
         task_feature_index = dataset.task_feature_index
-        task_values = dataset.X[..., dataset.task_feature_index].unique().long()
     ssd = none_throws(search_space_digest)
-    if (ssd.target_values is not None) and (
-        target_value := ssd.target_values.get(none_throws(task_feature_index))
-    ) is not None:
-        outcome_transform_options.setdefault("default_task_value", int(target_value))
     outcome_transform_options.setdefault("stratification_idx", task_feature_index)
-    outcome_transform_options.setdefault("observed_task_values", task_values)
     outcome_transform_options.setdefault(
         "all_task_values",
         torch.tensor(

ax/generators/torch/tests/test_outcome_transform_argparse.py

@@ -88,7 +88,7 @@ def test_argparse_stratified_standardize(self) -> None:
             dataset=mt_dataset,
             search_space_digest=ssd,
         )
-        options_b = {"stratification_idx": 2, "default_task_value": 4}
+        options_b = {"stratification_idx": 2}
         outcome_transform_kwargs_b = outcome_transform_argparse(
             StratifiedStandardize,
             dataset=mt_dataset,
@@ -97,27 +97,24 @@ def test_argparse_stratified_standardize(self) -> None:
         )
         expected_options_a = {
             "stratification_idx": 3,
-            "observed_task_values": torch.tensor([0, 1], dtype=torch.long),
             "all_task_values": torch.tensor([0, 1, 2], dtype=torch.long),
-            "default_task_value": 1,
         }
         expected_options_b = {
             "stratification_idx": 2,
-            "observed_task_values": torch.tensor([0, 1], dtype=torch.long),
             "all_task_values": torch.tensor([0, 1, 2], dtype=torch.long),
-            "default_task_value": 4,
         }
         for expected_options, actual_options in zip(
             (expected_options_a, expected_options_b),
             (outcome_transform_kwargs_a, outcome_transform_kwargs_b),
         ):
-            self.assertEqual(len(actual_options), 4)
-            for k in ("stratification_idx", "stratification_idx"):
-                self.assertEqual(actual_options[k], expected_options[k])
-            for k in ("observed_task_values", "all_task_values"):
-                self.assertTrue(
-                    torch.equal(
-                        actual_options[k],
-                        assert_is_instance(expected_options[k], Tensor),
-                    )
+            self.assertEqual(len(actual_options), 2)
+            self.assertEqual(
+                actual_options["stratification_idx"],
+                expected_options["stratification_idx"],
+            )
+            self.assertTrue(
+                torch.equal(
+                    actual_options["all_task_values"],
+                    assert_is_instance(expected_options["all_task_values"], Tensor),
                 )
+            )