Migrate transforms and utils to RangeParameter.step_size#5222
Open
saitcakmak wants to merge 4 commits into
Open
Migrate transforms and utils to RangeParameter.step_size#5222saitcakmak wants to merge 4 commits into
saitcakmak wants to merge 4 commits into
Conversation
Summary: The TL adapter utilities `get_joint_search_space`, `merge_dependents`, `merge_parameters` (in `ax/adapter/transfer_learning/utils.py`) and `get_mapped_parameter_names` (in `ax/adapter/transfer_learning/utils_torch.py`) were previously migrated out of `ax/fb/adapter/`, which left behind pure re-export shims at `ax/fb/adapter/utils.py` and `ax/fb/adapter/utils_torch.py`. The only remaining coverage for these functions lived in `ax/fb/adapter/tests/test_utils.py` and `test_utils_torch.py`, exercising the migrated code through those shims -- and the non-fb destination had no test coverage of its own. This moves both test files to `ax/adapter/transfer_learning/tests/`, switches their imports to the real non-fb modules (`ax.adapter.transfer_learning.utils`/`utils_torch`, `ax.core.auxiliary_source.AuxiliarySource`, `ax.utils.common.testutils.TestCase`), and removes the now-unused `get_unordered_choice`/`get_ordered_choice` helpers from `test_utils.py`. Since the two test files were the only callers of the `ax.fb.adapter.utils`/`utils_torch` shims repo-wide, the shims are deleted and the BUCK targets are updated accordingly: a new `test_utils` `python_unittest` is added under `ax/adapter/transfer_learning/BUCK`, and the old `test_utils` target, the orphaned `:utils` library, and the `utils_torch.py` src are removed from `ax/fb/adapter/BUCK`. The broader `ax.fb.core.auxiliary_source` shim is left in place; it still has many callers across admarket, pts, automl, storage, and docs, so cleaning it up is a separate effort. Differential Revision: D107429272
Summary: Adds a `step_size` arg to `RangeParameter` that snaps values to a grid anchored at `lower` (in `cast()`), for both FLOAT and INT parameters. This is the first diff in the step_size unification stack (see `ax/api/configs/RFC_step_size_unification.md`): `step_size` will subsume both the discrete-grid and limited-resolution (`digits`) use cases under one knob. In this diff `step_size` coexists with the existing `digits` arg (they are mutually exclusive at construction). Subsequent diffs in the stack migrate storage (JSON + SQA), transforms and utils, and the public API (`RangeParameterConfig`) to `step_size`, then deprecate `digits` in favor of it. Behavior: - `cast()` rounds `(value - lower) / step_size` to the nearest integer and returns `lower + n * step_size`. It does NOT clamp to `[lower, upper]`: an out-of-bounds input (e.g. a historical observation recorded outside the current bounds) snaps to the nearest grid point, which may itself be out of bounds. This mirrors the non-`step_size` `cast()`, which leaves out-of-bounds values in place rather than silently moving them into range — range validity is enforced by `validate()`, not `cast()`. - Both bounds must lie on the grid: `(upper - lower)` must be an integer multiple of `step_size` (within `EPS`). Off-grid bounds are rejected at construction. This guarantees `upper` is itself a feasible value, so a value near the upper bound snaps to `upper` rather than to a grid point short of it. - `step_size` must be strictly positive, and must be integer-valued for INT parameters. - `cardinality()` accounts for `step_size`: a grid-valued FLOAT reports the finite number of grid points instead of `inf`, and a grid-valued INT counts grid points rather than every integer in `[lower, upper]`. `step_size` defines a discrete grid but does not, by itself, force discrete acquisition optimization; how the optimizer treats the parameter depends on the grid cardinality and is determined at the generator level. Differential Revision: D107274057
Summary: Teaches both storage backends to persist and restore the new `step_size` field on `RangeParameter`. Second diff in the step_size unification stack (builds on the native `step_size` support added in D107274057). The underlying SQA column (`SQAParameter.step_size`) and the corresponding EntAE schema were added and landed separately (D107108021, D107112306), so this diff is pure encoder/decoder logic — it does not touch the schema. SQA: the encoder writes `parameter.step_size` to the column (and continues writing the legacy `digits` column for now; `digits` is dropped in a later diff). The decoder prefers `step_size` and falls back to the legacy `digits` column for rows written by older code, passing only one to the constructor (which rejects both being set; it converts `digits` to `step_size` internally). JSON: `range_parameter_to_dict` emits `step_size` alongside `digits`. There is no custom JSON decoder for `RangeParameter` — the registry maps it directly to the class, and `ax_class_from_json_dict` splats the serialized dict straight into the constructor. So `step_size` flows through automatically, and legacy blobs carrying only `digits` still decode via the constructor's back-compat handling. No decoder change is required for JSON. Differential Revision: D107282941
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
meta-cla
Bot
added
the
CLA Signed
|
@saitcakmak has exported this pull request. If you are a Meta employee, you can view the originating Diff in D107284896. |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #5222 +/- ##
==========================================
+ Coverage 96.50% 96.58% +0.07%
==========================================
Files 617 619 +2
Lines 69776 70113 +337
==========================================
+ Hits 67339 67720 +381
+ Misses 2437 2393 -44 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary:
Teaches the transform/util layer about step_size. Third diff in the step_size unification stack.
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