Updating SearchSpace._validate_derived_parameter to allow string and Boolean types for simple copies (#4851)

Summary:
Updating `SearchSpace._validate_derived_parameter` in light of the recent changes to `DerivedParameter`, enabling string and Boolean types for simple copy expressions.

This commit will resolve the remaining error in https://www.internalfb.com/ax/experiment/derived_parameter_string_bool/trials.

Pull Request resolved: #4851

Reviewed By: sunnyshen321

Differential Revision: D92164696

fbshipit-source-id: ebd8bdf74e76d77ff6d3ab8daa1778531ff29e63

Author: SebastianAment

ax/core/search_space.py

@@ -896,24 +896,49 @@ def validate_membership(self, parameters: TParameterization) -> None:
 
     def _validate_derived_parameter(self, parameter: DerivedParameter) -> None:
         is_int = parameter.parameter_type == ParameterType.INT
+        is_simple_copy = parameter._is_simple_copy
+        derived_is_numeric = parameter.parameter_type in (
+            ParameterType.INT,
+            ParameterType.FLOAT,
+        )
+
         for p_name in parameter.parameter_names_to_weights.keys():
             p = self._parameters.get(p_name)
             if p is None:
                 raise ValueError(
                     f"Parameter {p_name} is not in the search space, but is used in a "
                     "derived parameter."
                 )
-            if not p.is_numeric:
+
+            # For arithmetic expressions, source must be numeric
+            if not is_simple_copy and not p.is_numeric:
                 raise ValueError(
                     f"Parameter {p_name} is not a float or int, but is used in a "
-                    "derived parameter."
+                    "derived parameter whose expression is not a simple copy."
+                )
+
+            # For simple copies, validate type compatibility
+            # Valid: exact type match OR both numeric (int can promote to float)
+            if is_simple_copy:
+                types_compatible = parameter.parameter_type == p.parameter_type or (
+                    derived_is_numeric and p.is_numeric
                 )
-            elif is_int and p.parameter_type == ParameterType.FLOAT:
+                if not types_compatible:
+                    raise ValueError(
+                        f"Parameter {p_name} has type {p.parameter_type.name}, but the "
+                        f"derived parameter has type {parameter.parameter_type.name}. "
+                        "Simple copy derived parameters must have the same type as "
+                        "their source parameter."
+                    )
+
+            # Float source cannot be used with Int derived parameter
+            if is_int and p.parameter_type == ParameterType.FLOAT:
                 raise ValueError(
                     f"Parameter {p_name} is a float, but is used in a derived "
                     "parameter with int type."
                 )
-            elif isinstance(p, DerivedParameter):
+
+            if isinstance(p, DerivedParameter):
                 raise ValueError(
                     "Parameter cannot be derived from another derived parameter."
                 )

ax/core/tests/test_search_space.py

@@ -578,15 +578,42 @@ def test_validate_derived_parameter(self) -> None:
         ):
             self.ss1._validate_derived_parameter(parameter=self.invalid_derived_param)
 
-        # test with non-numeric param
+        # test with non-numeric param used in arithmetic expression
         derived_param = DerivedParameter(
-            name="z", parameter_type=ParameterType.FLOAT, expression_str="c"
+            name="z", parameter_type=ParameterType.FLOAT, expression_str="2.0 * c"
         )
         with self.assertRaisesRegex(
             ValueError,
             "Parameter c is not a float or int, but is used in a derived parameter.",
         ):
             self.ss1._validate_derived_parameter(parameter=derived_param)
+
+        # test simple copy type incompatibility: numeric derived from non-numeric source
+        # tests the unified type compatibility rule: types must match OR both numeric
+        derived_param = DerivedParameter(
+            name="z", parameter_type=ParameterType.FLOAT, expression_str="c"
+        )
+        with self.assertRaisesRegex(
+            ValueError,
+            "Parameter c has type STRING, but the derived parameter has type FLOAT. "
+            "Simple copy derived parameters must have the same type as their source "
+            "parameter.",
+        ):
+            self.ss1._validate_derived_parameter(parameter=derived_param)
+
+        # test simple copy type incompatibility: non-numeric derived from numeric source
+        # same validation rule as above, different type combination
+        derived_param = DerivedParameter(
+            name="z", parameter_type=ParameterType.STRING, expression_str="a"
+        )
+        with self.assertRaisesRegex(
+            ValueError,
+            "Parameter a has type FLOAT, but the derived parameter has type STRING. "
+            "Simple copy derived parameters must have the same type as their source "
+            "parameter.",
+        ):
+            self.ss1._validate_derived_parameter(parameter=derived_param)
+
         # test int derived param with float constituent param
         derived_param = DerivedParameter(
             name="z", parameter_type=ParameterType.INT, expression_str="a"
@@ -597,6 +624,16 @@ def test_validate_derived_parameter(self) -> None:
         ):
             self.ss1._validate_derived_parameter(parameter=derived_param)
 
+        # test int derived param with float constituent param (arithmetic expression)
+        derived_param = DerivedParameter(
+            name="z", parameter_type=ParameterType.INT, expression_str="2.0 * a"
+        )
+        with self.assertRaisesRegex(
+            ValueError,
+            "Parameter a is a float, but is used in a derived parameter with int type.",
+        ):
+            self.ss1._validate_derived_parameter(parameter=derived_param)
+
         # test derived param with constituent derived param
         derived_param = DerivedParameter(
             name="z", parameter_type=ParameterType.FLOAT, expression_str="h"
@@ -621,6 +658,37 @@ def test_validate_derived_parameter(self) -> None:
         ):
             self.ss1._validate_derived_parameter(parameter=derived_param)
 
+        # test simple copy of STRING parameter - should succeed
+        string_derived_param = DerivedParameter(
+            name="derived_c", parameter_type=ParameterType.STRING, expression_str="c"
+        )
+        # This should NOT raise - it's a valid simple copy
+        self.ss1._validate_derived_parameter(parameter=string_derived_param)
+
+        # test simple copy of BOOL parameter - should succeed
+        # Add a non-fixed BOOL parameter to the search space
+        bool_choice_param = ChoiceParameter(
+            name="bool_choice", parameter_type=ParameterType.BOOL, values=[True, False]
+        )
+        self.ss1.add_parameter(bool_choice_param)
+        bool_derived_param = DerivedParameter(
+            name="derived_bool",
+            parameter_type=ParameterType.BOOL,
+            expression_str="bool_choice",
+        )
+        # This should NOT raise - it's a valid simple copy
+        self.ss1._validate_derived_parameter(parameter=bool_derived_param)
+
+        # test simple copy INT to FLOAT promotion - should succeed
+        # INT can be promoted to FLOAT (e.g., 3 -> 3.0)
+        int_to_float_derived_param = DerivedParameter(
+            name="derived_f_as_float",
+            parameter_type=ParameterType.FLOAT,
+            expression_str="f",  # f is an INT parameter
+        )
+        # This should NOT raise - INT can be promoted to FLOAT
+        self.ss1._validate_derived_parameter(parameter=int_to_float_derived_param)
+
     def test_get_overlapping_parameters(self) -> None:
         with self.subTest("full_overlap"):
             range_param_1 = RangeParameter(