Clear digits in UnitX and Logit transforms before updating bounds (#5071)

Summary:
Pull Request resolved: #5071

When a RangeParameter has a `digits` value calibrated for its original scale (e.g., `digits=-3` to round to the nearest 1000 on a [5000, 500000] range), transforms that change the parameter's scale must clear `digits` before calling `update_range`. Otherwise, `RangeParameter.cast()` applies the original-scale rounding to the transformed bounds, which can collapse them (e.g., `round(1.0, -3) == 0.0`, making both bounds 0.0 and raising `UserInputError`).

The `Log` transform already handles this correctly (D66670173). This diff applies the same fix to `UnitX` and `Logit`, and documents the pattern in the base `Transform.transform_search_space` docstring.

Reviewed By: saitcakmak

Differential Revision: D97223066

fbshipit-source-id: d3c46b6963ad2ef594486cea97e4cb49e48bfbff

Author: ItsMrLin

ax/adapter/transforms/base.py

@@ -109,6 +109,14 @@ def transform_search_space(self, search_space: SearchSpace) -> SearchSpace:
         This is typically done in-place. This class implements the identity
         transform (does nothing).
 
+        NOTE for subclasses: If a transform changes the *scale* of a
+        RangeParameter (e.g., Log, UnitX, Logit), it must clear ``digits``
+        via ``p.set_digits(digits=None)`` before calling ``update_range``.
+        Otherwise, rounding calibrated for the original scale will corrupt
+        the transformed bounds (e.g., ``digits=-3`` rounds to the nearest
+        1000, which collapses [0, 1] to 0). The Cast transform re-applies
+        ``digits`` in the original space during untransform.
+
         Args:
             search_space: The search space
 

ax/adapter/transforms/logit.py

@@ -66,6 +66,10 @@ def transform_observation_features(
     def transform_search_space(self, search_space: SearchSpace) -> SearchSpace:
         for p_name, p in search_space.parameters.items():
             if p_name in self.transform_parameters and isinstance(p, RangeParameter):
+                # Don't round in logit space; digits will be re-applied in
+                # the original space by the Cast transform during untransform.
+                if p.digits is not None:
+                    p.set_digits(digits=None)
                 p.set_logit_scale(False).update_range(
                     lower=logit(p.lower).item(), upper=logit(p.upper).item()
                 )

ax/adapter/transforms/tests/test_logit_transform.py

@@ -122,6 +122,28 @@ def test_TransformSearchSpace(self) -> None:
         self.assertEqual(x_param.lower, logit(0.1))
         self.assertEqual(x_param.upper, logit(0.3))
 
+    def test_transform_search_space_clears_digits(self) -> None:
+        """Test that digits is cleared during transform to avoid rounding
+        in logit space."""
+        ss = SearchSpace(
+            parameters=[
+                RangeParameter(
+                    "x",
+                    lower=0.1,
+                    upper=0.9,
+                    parameter_type=ParameterType.FLOAT,
+                    logit_scale=True,
+                    digits=3,
+                ),
+            ]
+        )
+        t = Logit(search_space=ss)
+        ss = t.transform_search_space(ss)
+        x = assert_is_instance(ss.parameters["x"], RangeParameter)
+        self.assertIsNone(x.digits)
+        self.assertAlmostEqual(x.lower, logit(0.1))
+        self.assertAlmostEqual(x.upper, logit(0.9))
+
     def test_transform_experiment_data(self) -> None:
         parameterizations = [
             {"x": 0.2, "a": 1, "b": "a"},

ax/adapter/transforms/tests/test_unit_x_transform.py

@@ -133,6 +133,29 @@ def test_TransformSearchSpace(self) -> None:
             self.search_space_with_target.parameters["x"].target_value, 1.0
         )
 
+    def test_transform_search_space_clears_digits(self) -> None:
+        """Test that digits is cleared during transform to avoid rounding
+        in unit space. Regression test for a bug where digits=-3 (round to
+        nearest 1000) collapsed [0, 1] bounds to (0.0, 0.0)."""
+        ss = SearchSpace(
+            parameters=[
+                RangeParameter(
+                    "w",
+                    lower=5000.0,
+                    upper=500000.0,
+                    parameter_type=ParameterType.FLOAT,
+                    digits=-3,
+                ),
+            ]
+        )
+        t = UnitX(search_space=ss)
+        ss = t.transform_search_space(ss)
+        w = assert_is_instance(ss.parameters["w"], RangeParameter)
+        # digits must be cleared so rounding doesn't corrupt [0, 1] bounds.
+        self.assertIsNone(w.digits)
+        self.assertEqual(w.lower, 0.0)
+        self.assertEqual(w.upper, 1.0)
+
     def test_TransformNewSearchSpace(self) -> None:
         new_ss = SearchSpace(
             parameters=[

ax/adapter/transforms/unit_x.py

@@ -73,6 +73,10 @@ def transform_search_space(self, search_space: SearchSpace) -> SearchSpace:
             if (p_bounds := self.bounds.get(p_name)) is not None and isinstance(
                 p, RangeParameter
             ):
+                # Don't round in unit space; digits will be re-applied in
+                # the original space by the Cast transform during untransform.
+                if p.digits is not None:
+                    p.set_digits(digits=None)
                 p.update_range(
                     lower=self._normalize_value(value=p.lower, bounds=p_bounds),
                     upper=self._normalize_value(value=p.upper, bounds=p_bounds),