Migrate transforms and utils to RangeParameter.step_size

Summary:
Teaches the transform/util layer about step_size. Third diff in the step_size unification stack.

- Cast: snaps RangeParameter values to the grid via parameter.cast() on both the observation-features path (already calls cast) and the experiment_data dataframe path (replacing the .round(digits) call), for params with digits OR step_size set.
- Log/Logit/UnitX: clear step_size before rescaling (in addition to clearing digits, until digits is fully removed in a later diff); step_size is re-applied in the original space by Cast on untransform.
- int_to_float: does not forward step_size to the FLOAT surrogate (anchor would be misaligned); the original INT param's snapping is re-applied by Cast. See TODO.
- map_key_to_float, transfer_learning merge_parameters, service instantiation: forward/read step_size, preferring it over digits.
- service instantiation: add step_size to EXPECTED_KEYS_IN_PARAM_REPR. The RangeParameter construction path already read representation["step_size"], but the key was not in the recognized-keys set, so any parameter representation passing step_size would have been rejected with an "Unexpected keys" error. This makes the step_size representation path actually usable.
- core_stubs: add get_range_parameter_with_step_size helper.

Test coverage added for the migrated paths: Log/Logit transform_search_space clears step_size (mirroring the existing UnitX and clears_digits tests), MapKeyToFloat forwards step_size from config, merge_parameters forwards step_size/digits from p1, and parameter_from_json accepts step_size and legacy digits for range parameters.

Differential Revision: D107284896

Author: saitcakmak

ax/adapter/transfer_learning/tests/test_utils.py

@@ -254,6 +254,41 @@ def test_merge_parameters(self) -> None:
                 name="p", parameter_type=ParameterType.INT, lower=0, upper=3
             ),
         )
+        # The grid spec (step_size, or legacy digits) is forwarded from p1, even
+        # when the merged bounds are widened by p2 (as long as they stay on p1's
+        # grid).
+        p_range_step = RangeParameter(
+            name="p",
+            parameter_type=ParameterType.FLOAT,
+            lower=0.0,
+            upper=1.0,
+            step_size=0.1,
+        )
+        p_range_wide = RangeParameter(
+            name="p", parameter_type=ParameterType.FLOAT, lower=0.0, upper=2.0
+        )
+        merged_step = assert_is_instance(
+            merge_parameters(p1=p_range_step, p2=p_range_wide, reverse_param_config={}),
+            RangeParameter,
+        )
+        self.assertEqual(merged_step.upper, 2.0)
+        self.assertEqual(merged_step.step_size, 0.1)
+        self.assertIsNone(merged_step.digits)
+        p_range_digits = RangeParameter(
+            name="p",
+            parameter_type=ParameterType.FLOAT,
+            lower=0.0,
+            upper=1.0,
+            digits=2,
+        )
+        merged_digits = assert_is_instance(
+            merge_parameters(
+                p1=p_range_digits, p2=p_range_wide, reverse_param_config={}
+            ),
+            RangeParameter,
+        )
+        self.assertEqual(merged_digits.digits, 2)
+        self.assertIsNone(merged_digits.step_size)
         p_choice_1 = ChoiceParameter(
             name="p",
             parameter_type=ParameterType.STRING,

ax/adapter/transfer_learning/utils.py

@@ -111,6 +111,7 @@ def merge_parameters(
             log_scale=p1.log_scale,
             logit_scale=p1.logit_scale,
             digits=p1.digits,
+            step_size=p1.step_size,
             is_fidelity=p1.is_fidelity,
             target_value=p1.target_value,
         )

ax/adapter/transforms/base.py

@@ -94,12 +94,12 @@ def transform_search_space(self, search_space: SearchSpace) -> SearchSpace:
         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.
+        RangeParameter (e.g., Log, UnitX, Logit), it must clear ``step_size``
+        via ``p.set_step_size(step_size=None)`` before calling ``update_range``.
+        Otherwise, snapping calibrated for the original scale will corrupt the
+        transformed bounds (a grid spacing meaningful in the original space is
+        meaningless after a non-linear rescale). The Cast transform re-applies
+        ``step_size`` in the original space during untransform.
 
         Args:
             search_space: The search space

ax/adapter/transforms/cast.py

@@ -25,7 +25,7 @@
 from ax.exceptions.core import UserInputError
 from ax.generators.types import TConfig
 from ax.utils.common.constants import Keys
-from pandas import DataFrame
+from pandas import DataFrame, Series
 from pyre_extensions import assert_is_instance, none_throws
 
 if TYPE_CHECKING:
@@ -314,11 +314,38 @@ def transform_experiment_data(
             for p, param in self.search_space.parameters.items()
         }
         arm_data = arm_data.astype(dtype=column_to_type)
-        # Round to digits if any parameter specifies it.
+        # Snap to the parameter's grid (digits or step_size) if specified.
+        # These mirror ``RangeParameter.cast``'s rounding logic, but are applied
+        # in a vectorized manner over the whole column rather than via a per-row
+        # ``Series.apply`` (which calls ``parameter.cast`` once per element and is
+        # slow for large DataFrames). NaN / ``<NA>`` values (added for missing
+        # columns during the ``reindex`` above) propagate through ``round`` and
+        # the arithmetic, matching the previous ``value if value is None`` guard.
         for p_name in parameter_names:
             parameter = self.search_space.parameters[p_name]
-            if isinstance(parameter, RangeParameter) and parameter.digits is not None:
-                arm_data[p_name] = arm_data[p_name].round(parameter.digits)
+            if not isinstance(parameter, RangeParameter):
+                continue
+            column: Series = arm_data[p_name]
+            if (
+                parameter.parameter_type is ParameterType.FLOAT
+                and parameter.digits is not None
+            ):
+                # ``Series.round`` uses round-half-to-even, same as Python's
+                # built-in ``round`` used in ``RangeParameter.cast``.
+                arm_data[p_name] = column.round(parameter.digits)
+            elif parameter.step_size is not None:
+                # Snap to the grid ``{lower + k * step_size : k in Z}`` by
+                # rounding ``(value - lower) / step_size`` to the nearest integer.
+                lower = float(parameter.lower)
+                step_size = none_throws(parameter.step_size)
+                steps: Series = column.sub(lower).div(step_size).round()
+                snapped: Series = steps.mul(step_size).add(lower)
+                if parameter.parameter_type is ParameterType.INT:
+                    # Preserve the nullable ``Int64`` dtype so reindex-added
+                    # ``<NA>`` values survive the cast.
+                    arm_data[p_name] = snapped.round().astype("Int64")
+                else:
+                    arm_data[p_name] = snapped
 
         return ExperimentData(arm_data=arm_data, observation_data=observation_data)
 

ax/adapter/transforms/log.py

@@ -79,7 +79,12 @@ def transform_search_space(self, search_space: SearchSpace) -> SearchSpace:
                     isinstance(p, RangeParameter)
                     and p.parameter_type == ParameterType.FLOAT
                 ):
-                    # Don't round in log space
+                    # Don't snap/round in log space; step_size (or legacy
+                    # digits) will be re-applied in the original space by the
+                    # Cast transform during untransform. Both are cleared until
+                    # digits is fully removed (see step_size unification RFC).
+                    if p.step_size is not None:
+                        p.set_step_size(step_size=None)
                     if p.digits is not None:
                         p.set_digits(digits=None)
                     p.set_log_scale(False).update_range(

ax/adapter/transforms/logit.py

@@ -66,8 +66,12 @@ 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.
+                # Don't snap/round in logit space; step_size (or legacy digits)
+                # will be re-applied in the original space by the Cast transform
+                # during untransform. Both are cleared until digits is fully
+                # removed (see step_size unification RFC).
+                if p.step_size is not None:
+                    p.set_step_size(step_size=None)
                 if p.digits is not None:
                     p.set_digits(digits=None)
                 p.set_logit_scale(False).update_range(

ax/adapter/transforms/map_key_to_float.py

@@ -127,13 +127,18 @@ def __init__(
                 return
 
             p_config = self.parameters[MAP_KEY]
+            # Prefer ``step_size``; fall back to legacy ``digits``. Only one may
+            # be passed to the constructor (it rejects both being set).
+            step_size = p_config.get("step_size", None)
+            digits = None if step_size is not None else p_config.get("digits", None)
             self._parameter_list.append(
                 RangeParameter(
                     name=MAP_KEY,
                     parameter_type=ParameterType.FLOAT,
                     lower=p_config.get("lower", min(values)),
                     upper=p_config.get("upper", max(values)),
-                    digits=p_config.get("digits", None),
+                    digits=digits,
+                    step_size=step_size,
                     is_fidelity=p_config.get("is_fidelity", False),
                     target_value=p_config.get("target_value", None),
                 )

ax/adapter/transforms/tests/test_cast_transform.py

@@ -323,6 +323,77 @@ def test_cast_parameter_type_and_none(self) -> None:
         ]
         self.assertEqual(tf_observations, expected)
 
+    def test_cast_step_size_observation_features(self) -> None:
+        # Cast snaps RangeParameter values to step_size on (un)transform, just
+        # as it rounds to digits.
+        search_space = SearchSpace(
+            parameters=[
+                RangeParameter(
+                    name="range",
+                    parameter_type=ParameterType.FLOAT,
+                    lower=0.0,
+                    upper=1.0,
+                    step_size=0.1,
+                ),
+            ]
+        )
+        t = Cast(search_space=search_space)
+        obs_features = [
+            ObservationFeatures(parameters={"range": 0.12}),
+            ObservationFeatures(parameters={"range": 0.36}),
+        ]
+        tf_obs_features = t.transform_observation_features(
+            observation_features=obs_features
+        )
+        self.assertAlmostEqual(
+            float(none_throws(tf_obs_features[0].parameters["range"])), 0.1
+        )
+        self.assertAlmostEqual(
+            float(none_throws(tf_obs_features[1].parameters["range"])), 0.4
+        )
+
+    def test_transform_experiment_data_step_size(self) -> None:
+        # The experiment_data dataframe path snaps RangeParameter values to
+        # step_size, for both FLOAT and INT parameters. The INT parameter also
+        # checks that the snapped column keeps the nullable Int64 dtype.
+        experiment = get_experiment_with_observations(
+            observations=[[0.0], [1.0]],
+            search_space=SearchSpace(
+                parameters=[
+                    RangeParameter(
+                        name="x",
+                        parameter_type=ParameterType.FLOAT,
+                        lower=0.0,
+                        upper=1.0,
+                        step_size=0.1,
+                    ),
+                    RangeParameter(
+                        name="y",
+                        parameter_type=ParameterType.INT,
+                        lower=0,
+                        upper=10,
+                        step_size=2,
+                    ),
+                ]
+            ),
+            parameterizations=[
+                {"x": 0.12, "y": 3},
+                {"x": 0.36, "y": 7},
+            ],
+        )
+        experiment_data = extract_experiment_data(
+            experiment=experiment, data_loader_config=DataLoaderConfig()
+        )
+        transformed = Cast(
+            search_space=experiment.search_space
+        ).transform_experiment_data(experiment_data=deepcopy(experiment_data))
+        self.assertAlmostEqual(transformed.arm_data["x"].iloc[0], 0.1)
+        self.assertAlmostEqual(transformed.arm_data["x"].iloc[1], 0.4)
+        # 3 snaps to 4 (round half to even: 1.5 -> 2 steps), 7 snaps to 8.
+        self.assertEqual(transformed.arm_data["y"].iloc[0], 4)
+        self.assertEqual(transformed.arm_data["y"].iloc[1], 8)
+        self.assertEqual(transformed.arm_data["y"].dtype, "Int64")
+
     def test_transform_experiment_data_flatten(self) -> None:
         # Tests for flattening of hierarchical parameterizations.
         columns = [

ax/adapter/transforms/tests/test_log_transform.py

@@ -117,12 +117,36 @@ def test_TransformSearchSpace(self) -> None:
         ss2 = deepcopy(self.search_space)
         ss2 = self.t.transform_search_space(ss2)
 
-        # Test float log-scale parameter transformation
+        # Test float log-scale parameter transformation. The grid (legacy
+        # ``digits`` here; ``step_size`` covered below) must be cleared during
+        # the transform -- a grid meaningful in the original space is
+        # meaningless after a log10 rescale, and Cast re-applies it in the
+        # original space on untransform.
         param_x = assert_is_instance(ss2.parameters["x"], RangeParameter)
         self.assertEqual(param_x.lower, math.log10(1))
         self.assertEqual(param_x.upper, math.log10(3))
         self.assertIsNone(param_x.digits)
 
+        # Same clearing behavior for ``step_size`` (mutually exclusive with
+        # ``digits``, so it needs its own parameter).
+        ss_step = SearchSpace(
+            parameters=[
+                RangeParameter(
+                    "x",
+                    lower=1.0,
+                    upper=1000.0,
+                    parameter_type=ParameterType.FLOAT,
+                    log_scale=True,
+                    step_size=1.0,
+                ),
+            ]
+        )
+        ss_step = Log(search_space=ss_step).transform_search_space(ss_step)
+        param_x_step = assert_is_instance(ss_step.parameters["x"], RangeParameter)
+        self.assertIsNone(param_x_step.step_size)
+        self.assertEqual(param_x_step.lower, math.log10(1.0))
+        self.assertEqual(param_x_step.upper, math.log10(1000.0))
+
         # Test integer log-scale parameter transformation (converted to ChoiceParameter)
         param_y = assert_is_instance(ss2.parameters["y"], ChoiceParameter)
         self.assertEqual(param_y.parameter_type, ParameterType.FLOAT)

ax/adapter/transforms/tests/test_logit_transform.py

@@ -122,27 +122,38 @@ 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_search_space_clears_grid(self) -> None:
+        """The grid (legacy ``digits`` or ``step_size``, which are mutually
+        exclusive) must be cleared during the transform to avoid rounding /
+        snapping in logit space; it is re-applied in the original space by Cast
+        on untransform."""
+        grid_params = [
+            RangeParameter(
+                "x",
+                lower=0.1,
+                upper=0.9,
+                parameter_type=ParameterType.FLOAT,
+                logit_scale=True,
+                digits=3,
+            ),
+            RangeParameter(
+                "x",
+                lower=0.1,
+                upper=0.9,
+                parameter_type=ParameterType.FLOAT,
+                logit_scale=True,
+                step_size=0.1,
+            ),
+        ]
+        for param in grid_params:
+            with self.subTest(param=param):
+                ss = SearchSpace(parameters=[param])
+                ss = Logit(search_space=ss).transform_search_space(ss)
+                x = assert_is_instance(ss.parameters["x"], RangeParameter)
+                self.assertIsNone(x.digits)
+                self.assertIsNone(x.step_size)
+                self.assertAlmostEqual(x.lower, logit(0.1))
+                self.assertAlmostEqual(x.upper, logit(0.9))
 
     def test_transform_experiment_data(self) -> None:
         parameterizations = [