Option to not include range bounds in SearchSpace.check_membership (facebook#4962)

Summary: Pull Request resolved: facebook#4962

Reviewed By: saitcakmak

Differential Revision: D94707497

fbshipit-source-id: 9f827e607939a282269d4793e51eb62cb3b282f9

Author: bletham

ax/analysis/utils.py

@@ -377,6 +377,7 @@ def _prepare_modeled_arm_data(
             parameterization=arm.parameters,
             raise_error=False,
             check_all_parameters_present=True,
+            check_range_bounds=False,
         ):
             predictable_pairs.append((trial_index, arm))
         else:

ax/core/search_space.py

@@ -341,6 +341,7 @@ def check_membership(
         parameterization: Mapping[str, TParamValue],
         raise_error: bool = False,
         check_all_parameters_present: bool = True,
+        check_range_bounds: bool = True,
     ) -> bool:
         """Whether the given parameterization belongs in the search space.
 
@@ -354,6 +355,10 @@ def check_membership(
                 with detailed explanation of why.
             check_all_parameters_present: Ensure that parameterization specifies
                 values for all parameters as expected by the search space.
+            check_range_bounds: If False, only check that values for
+                RangeParameters have the correct type, without enforcing
+                the parameter bounds. Other parameter types (ChoiceParameter,
+                FixedParameter, DerivedParameter) are still fully validated.
 
         Returns:
             Whether the parameterization is contained in the search space.
@@ -366,12 +371,16 @@ def check_membership(
 
         for name, value in parameterization.items():
             p = self.parameters[name]
-            kwargs = (
-                {"parameters": parameterization}
-                if isinstance(p, DerivedParameter)
-                else {}
-            )
-            if not p.validate(value=value, raises=False, **kwargs):
+            if not check_range_bounds and isinstance(p, RangeParameter):
+                is_valid = p.is_valid_type(value)
+            else:
+                kwargs = (
+                    {"parameters": parameterization}
+                    if isinstance(p, DerivedParameter)
+                    else {}
+                )
+                is_valid = p.validate(value=value, raises=False, **kwargs)
+            if not is_valid:
                 if raise_error:
                     raise ValueError(
                         f"{value} is not a valid value for "

ax/core/tests/test_search_space.py

@@ -303,6 +303,47 @@ def test_CheckMembership(self) -> None:
         with self.assertRaises(ValueError):
             self.ss2.check_membership(p_dict, raise_error=True)
 
+    def test_CheckMembershipSkipRangeBounds(self) -> None:
+        ss = SearchSpace(
+            parameters=[
+                RangeParameter(
+                    name="x",
+                    parameter_type=ParameterType.FLOAT,
+                    lower=0.0,
+                    upper=1.0,
+                ),
+                ChoiceParameter(
+                    name="c",
+                    parameter_type=ParameterType.STRING,
+                    values=["a", "b"],
+                ),
+            ]
+        )
+
+        # Value out of range: default check rejects it
+        p_dict = {"x": 5.0, "c": "a"}
+        self.assertFalse(ss.check_membership(p_dict))
+        self.assertFalse(ss.check_membership(p_dict, check_range_bounds=True))
+
+        # With check_range_bounds=False, out-of-range value is accepted
+        self.assertTrue(ss.check_membership(p_dict, check_range_bounds=False))
+
+        # Wrong type for range parameter is still rejected
+        p_dict_wrong_type = {"x": "not_a_number", "c": "a"}
+        self.assertFalse(
+            ss.check_membership(p_dict_wrong_type, check_range_bounds=False)
+        )
+        with self.assertRaises(ValueError):
+            ss.check_membership(
+                p_dict_wrong_type, check_range_bounds=False, raise_error=True
+            )
+
+        # Invalid choice value is still rejected
+        p_dict_bad_choice = {"x": 5.0, "c": "invalid"}
+        self.assertFalse(
+            ss.check_membership(p_dict_bad_choice, check_range_bounds=False)
+        )
+
     def test_check_membership_df(self) -> None:
         """Test vectorized membership check on DataFrames."""
         # Create test DataFrame with valid and invalid rows