diff --git a/ax/adapter/transfer_learning/tests/test_utils.py b/ax/adapter/transfer_learning/tests/test_utils.py index ccc644778f7..bd3417cb87c 100644 --- a/ax/adapter/transfer_learning/tests/test_utils.py +++ b/ax/adapter/transfer_learning/tests/test_utils.py @@ -26,6 +26,36 @@ class AxFbCoreUtilsTest(TestCase): + def test_merge_choice_parameters_preserves_log_scale(self) -> None: + for log_scale, extra_value in [(True, 3), (False, 4)]: + choice = ChoiceParameter( + "x", + parameter_type=ParameterType.INT, + values=[1, 2], + is_ordered=True, + sort_values=True, + log_scale=log_scale, + ) + other_choice = ChoiceParameter( + "x", + parameter_type=ParameterType.INT, + values=[2, extra_value], + is_ordered=True, + sort_values=True, + log_scale=not log_scale, + ) + fixed = FixedParameter( + "x", parameter_type=ParameterType.INT, value=extra_value + ) + for p1, p2 in [(choice, other_choice), (choice, fixed), (fixed, choice)]: + with self.subTest(log_scale=log_scale, p1=p1, p2=p2): + merged = assert_is_instance( + merge_parameters(p1, p2, reverse_param_config={}), + ChoiceParameter, + ) + self.assertEqual(merged.values, [1, 2, extra_value]) + self.assertEqual(merged.log_scale, log_scale) + def test_get_joint_search_space(self) -> None: parameters: list[Parameter] = [ RangeParameter(f"x{i}", parameter_type=ParameterType.INT, lower=0, upper=5) diff --git a/ax/adapter/transfer_learning/utils.py b/ax/adapter/transfer_learning/utils.py index a20fe584b86..02d6abf9c24 100644 --- a/ax/adapter/transfer_learning/utils.py +++ b/ax/adapter/transfer_learning/utils.py @@ -165,6 +165,7 @@ def merge_parameters( is_fidelity=choice_param.is_fidelity, target_value=choice_param.target_value, sort_values=choice_param.sort_values, + log_scale=choice_param.log_scale, dependents=merge_dependents( p1=p1, p2=p2, @@ -181,6 +182,7 @@ def merge_parameters( is_fidelity=p1.is_fidelity, target_value=p1.target_value, sort_values=p1.sort_values, + log_scale=p1.log_scale, dependents=merge_dependents( p1=p1, p2=p2, reverse_param_config=reverse_param_config ),