Fix dependent handling in Log transform

ax/adapter/transforms/log.py

@@ -94,6 +94,7 @@ def transform_search_space(self, search_space: SearchSpace) -> SearchSpace:
                 ) or isinstance(p, ChoiceParameter):
                     # Handle both int RangeParameter and ChoiceParameter
                     # by converting to log-transformed ChoiceParameter
+                    dependents: dict[TParamValue, list[str]] | None = None
                     if isinstance(p, RangeParameter):
                         lower = assert_is_instance(p.lower, int)
                         upper = assert_is_instance(p.upper, int)
@@ -104,6 +105,7 @@ def transform_search_space(self, search_space: SearchSpace) -> SearchSpace:
                         values = p.values
                         is_ordered = p.is_ordered
                         sort_values = p.sort_values
+                        dependents = p._dependents
 
                     # Apply log10 transformation
                     transformed_values = [
@@ -115,6 +117,11 @@ def transform_search_space(self, search_space: SearchSpace) -> SearchSpace:
                         target_value = math.log10(
                             assert_is_instance_of_tuple(target_value, (float, int))
                         )
+                    if dependents is not None:
+                        dependents = {
+                            math.log10(assert_is_instance_of_tuple(k, (float, int))): v
+                            for k, v in dependents.items()
+                        }
 
                     # Create new ChoiceParameter with transformed values.
                     choice_param = ChoiceParameter(
@@ -126,6 +133,7 @@ def transform_search_space(self, search_space: SearchSpace) -> SearchSpace:
                         target_value=target_value,
                         sort_values=sort_values,
                         log_scale=False,
+                        dependents=dependents,
                         bypass_cardinality_check=True,
                     )
 

ax/adapter/transforms/task_encode.py

@@ -10,12 +10,8 @@
 
 from typing import TYPE_CHECKING
 
-from ax.adapter.data_utils import ExperimentData
 from ax.adapter.transforms.choice_encode import ChoiceToNumericChoice
-from ax.adapter.transforms.utils import construct_new_search_space
-from ax.core.parameter import ChoiceParameter, Parameter, ParameterType
-from ax.core.search_space import SearchSpace
-from ax.generators.types import TConfig
+from ax.core.parameter import ChoiceParameter, Parameter
 
 if TYPE_CHECKING:
     # import as module to make sphinx-autodoc-typehints happy
@@ -34,27 +30,6 @@ class TaskChoiceToIntTaskChoice(ChoiceToNumericChoice):
     Transform is done in-place.
     """
 
-    def __init__(
-        self,
-        search_space: SearchSpace,
-        experiment_data: ExperimentData | None = None,
-        adapter: adapter_module.base.Adapter | None = None,
-        config: TConfig | None = None,
-    ) -> None:
-        super().__init__(
-            search_space=search_space,
-            experiment_data=experiment_data,
-            adapter=adapter,
-            config=config,
-        )
-        self.target_values: dict[str, int | None] = {
-            p.name: self.encoded_parameters[p.name][p.target_value]
-            if p.target_value is not None
-            else None
-            for p in search_space.parameters.values()
-            if self._should_encode(p=p)
-        }
-
     def _should_encode(self, p: Parameter) -> bool:
         """Check if a parameter should be encoded.
         Encodes task choice parameters.
@@ -68,34 +43,3 @@ def _should_encode(self, p: Parameter) -> bool:
                 )
             return True
         return False
-
-    def transform_search_space(self, search_space: SearchSpace) -> SearchSpace:
-        transformed_parameters: dict[str, Parameter] = {}
-        for p_name, p in search_space.parameters.items():
-            if p_name in self.encoded_parameters and isinstance(p, ChoiceParameter):
-                if p.is_fidelity:
-                    raise ValueError(
-                        f"Cannot choice-encode fidelity parameter {p_name}."
-                    )
-                # Choice(|K|) => Choice(0, K-1, is_task=True)
-                transformed_parameters[p_name] = ChoiceParameter(
-                    name=p_name,
-                    parameter_type=ParameterType.INT,
-                    values=list(range(len(p.values))),
-                    is_ordered=p.is_ordered,
-                    is_task=True,
-                    sort_values=True,
-                    target_value=self.target_values[p_name],
-                )
-            else:
-                transformed_parameters[p.name] = p
-        return construct_new_search_space(
-            search_space=search_space,
-            parameters=list(transformed_parameters.values()),
-            parameter_constraints=[
-                pc.clone_with_transformed_parameters(
-                    transformed_parameters=transformed_parameters
-                )
-                for pc in search_space.parameter_constraints
-            ],
-        )

ax/adapter/transforms/tests/test_log_transform.py

@@ -208,6 +208,12 @@ def test_log_scale_choice_parameter(self) -> None:
                     log_scale=True,
                     is_fidelity=True,
                     target_value=32,
+                    dependents={4: ["t"]},
+                ),
+                ChoiceParameter(
+                    "t",
+                    parameter_type=ParameterType.INT,
+                    values=[1, 2, 3],
                 ),
             ]
         )
@@ -262,6 +268,10 @@ def test_log_scale_choice_parameter(self) -> None:
         self.assertFalse(param_w.log_scale)
         self.assertTrue(param_w.is_fidelity)
         self.assertEqual(param_w.target_value, math.log10(32))
+        self.assertEqual(param_w.dependents, {math.log10(4): ["t"]})
+
+        # Verify that `t` is not transformed.
+        self.assertEqual(ss2.parameters["t"], search_space.parameters["t"])
 
     @mock_botorch_optimize
     def test_log_scale_choice_with_adapter(self) -> None:

ax/adapter/transforms/tests/test_task_encode_transform.py

@@ -35,6 +35,7 @@ def setUp(self) -> None:
                     values=["online", "offline"],
                     is_task=True,
                     target_value="online",
+                    dependents={"online": ["b"]},
                 ),
             ]
         )
@@ -75,6 +76,7 @@ def test_TransformSearchSpace(self) -> None:
         # pyre-fixme[16]: `Parameter` has no attribute `values`.
         self.assertEqual(ss2.parameters["c"].values, [0, 1])
         self.assertEqual(ss2.parameters["c"].target_value, 0)
+        self.assertEqual(ss2.parameters["c"].dependents, {0: ["b"]})
 
         # Test error if there are fidelities
         ss3 = SearchSpace(