Raise errors for unsupported hierarchical search space functionalities (#4374)

Summary:

Raise errors for unsupported HSS functionalities.

Differential Revision: D83626637

Author: kayween

ax/adapter/transforms/choice_encode.py

@@ -13,7 +13,7 @@
 from ax.adapter.transforms.utils import ClosestLookupDict, construct_new_search_space
 from ax.core.observation import ObservationFeatures
 from ax.core.parameter import ChoiceParameter, Parameter, ParameterType, RangeParameter
-from ax.core.search_space import SearchSpace
+from ax.core.search_space import HierarchicalSearchSpace, SearchSpace
 from ax.core.types import TParamValue
 from ax.generators.types import TConfig
 
@@ -69,6 +69,15 @@ def __init__(
                     zip(transformed_values, p.values)
                 )
 
+        if isinstance(search_space, HierarchicalSearchSpace):
+            for p_name, p in search_space.parameters.items():
+                if p.is_hierarchical and p.parameter_type == ParameterType.FLOAT:
+                    raise RuntimeError(
+                        f"{p_name} is a float hierarchical parameter. However, "
+                        "hierarchical parameters have to be integer-valued after Ax "
+                        f"transformations, but {self.__class__.__name__} would skip it."
+                    )
+
     def transform_observation_features(
         self, observation_features: list[ObservationFeatures]
     ) -> list[ObservationFeatures]:

ax/adapter/transforms/one_hot.py

@@ -16,7 +16,7 @@
 from ax.adapter.transforms.utils import construct_new_search_space
 from ax.core.observation import ObservationFeatures
 from ax.core.parameter import ChoiceParameter, Parameter, ParameterType, RangeParameter
-from ax.core.search_space import SearchSpace
+from ax.core.search_space import HierarchicalSearchSpace, SearchSpace
 from ax.core.types import TParameterization, TParamValue
 from ax.generators.types import TConfig
 from pyre_extensions import assert_is_instance
@@ -119,6 +119,14 @@ def __init__(
                         f"{p.name}{OH_PARAM_INFIX}_{i}" for i in range(encoded_len)
                     ]
 
+        if isinstance(search_space, HierarchicalSearchSpace):
+            for p_name, p in search_space.parameters.items():
+                if p.is_hierarchical and p_name in self.encoded_parameters:
+                    raise RuntimeError(
+                        f"Attempt to one-hot encode a hierarchical parameter {p_name}. "
+                        "This is not supported yet."
+                    )
+
     def transform_observation_features(
         self, observation_features: list[ObservationFeatures]
     ) -> list[ObservationFeatures]:

ax/adapter/transforms/tests/test_choice_encode_transform.py

@@ -97,6 +97,28 @@ def setUp(self) -> None:
     def test_init(self) -> None:
         self.assertEqual(list(self.t.encoded_parameters.keys()), ["d", "e"])
 
+        hierarchical_search_space = HierarchicalSearchSpace(
+            parameters=[
+                ChoiceParameter(
+                    name="x0",
+                    parameter_type=ParameterType.FLOAT,
+                    values=[1.0, 2.0],
+                    dependents={1.0: ["x1"], 2.0: ["x2"]},
+                ),
+                RangeParameter(
+                    name="x1", lower=0.0, upper=1.0, parameter_type=ParameterType.FLOAT
+                ),
+                RangeParameter(
+                    name="x2", lower=0.0, upper=1.0, parameter_type=ParameterType.FLOAT
+                ),
+            ]
+        )
+        with self.assertRaisesRegex(
+            expected_exception=RuntimeError,
+            expected_regex=("x0 is a float hierarchical parameter*"),
+        ):
+            ChoiceToNumericChoice(search_space=hierarchical_search_space.clone())
+
     def test_transform_observation_features(self) -> None:
         observation_features = self.observation_features
         obs_ft2 = deepcopy(observation_features)

ax/adapter/transforms/tests/test_one_hot_transform.py

@@ -16,7 +16,7 @@
 from ax.core.observation import ObservationFeatures
 from ax.core.parameter import ChoiceParameter, ParameterType, RangeParameter
 from ax.core.parameter_constraint import ParameterConstraint
-from ax.core.search_space import RobustSearchSpace, SearchSpace
+from ax.core.search_space import HierarchicalSearchSpace, RobustSearchSpace, SearchSpace
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.core_stubs import (
     get_experiment_with_observations,
@@ -80,6 +80,35 @@ def test_Init(self) -> None:
         self.assertEqual(list(self.t.encoded_parameters.keys()), ["b"])
         self.assertEqual(list(self.t2.encoded_parameters.keys()), ["b"])
 
+        hierarchical_search_space = HierarchicalSearchSpace(
+            parameters=[
+                ChoiceParameter(
+                    name="x0",
+                    parameter_type=ParameterType.STRING,
+                    values=["a", "b", "c"],
+                    is_ordered=False,
+                    dependents={"a": ["x1"], "b": ["x2"], "c": ["x3"]},
+                ),
+                RangeParameter(
+                    name="x1", lower=0.0, upper=1.0, parameter_type=ParameterType.FLOAT
+                ),
+                RangeParameter(
+                    name="x2", lower=0.0, upper=1.0, parameter_type=ParameterType.FLOAT
+                ),
+                RangeParameter(
+                    name="x3", lower=0.0, upper=1.0, parameter_type=ParameterType.FLOAT
+                ),
+            ]
+        )
+        with self.assertRaisesRegex(
+            expected_exception=RuntimeError,
+            expected_regex=(
+                "Attempt to one-hot encode a hierarchical parameter x0. "
+                "This is not supported yet."
+            ),
+        ):
+            OneHot(search_space=hierarchical_search_space.clone())
+
     def test_TransformObservationFeatures(self) -> None:
         observation_features = [self.observation_features]
         obs_ft2 = deepcopy(observation_features)

ax/core/parameter.py

@@ -724,6 +724,13 @@ def __init__(
                         f"Value {value} in `dependents` "
                         f"argument is not among the parameter values: {self.values}."
                     )
+            if len(dependents) >= 3:
+                raise NotImplementedError(
+                    f"Found {len(dependents)} differet values in the hierarchical "
+                    f"parameter {name}. However, as of now, only <= 2 different values "
+                    "are supported for hierarchical parameters."
+                )
+
         # NOTE: We don't need to check that dependent parameters actually exist as
         # that is done in `HierarchicalSearchSpace` constructor.
         self._dependents = dependents

ax/core/tests/test_parameter.py

@@ -479,6 +479,24 @@ def test_Hierarchical(self) -> None:
                 dependents={"c": "other_param"},
             )
 
+        # Test case that raises an error when there are more than 2 values in a
+        # hierarchical choice parameter. This test case should be removed once we add
+        # support for it.
+        with self.assertRaisesRegex(
+            NotImplementedError,
+            "Found 3 differet values in the hierarchical parameter x.*",
+        ):
+            ChoiceParameter(
+                name="x",
+                parameter_type=ParameterType.STRING,
+                values=["a", "b", "c"],
+                dependents={
+                    "a": ["1st_child"],
+                    "b": ["2nd_child"],
+                    "c": ["3rd_child"],
+                },
+            )
+
     def test_available_flags(self) -> None:
         choice_flags = [
             "is_fidelity",