From 0df11ccb01f4b2078869fa11d75abc49177e9c94 Mon Sep 17 00:00:00 2001
From: Louis Tiao <tiao@meta.com>
Date: Wed, 21 Oct 2026 08:29:56 -0700
Subject: [PATCH] Preserve log scaling when merging choice parameters

Summary:
# Why

Transfer learning combines parameter definitions from the target and auxiliary experiments. Rebuilding a choice parameter can silently replace its logarithmic or linear scale with a default inferred from the combined values. Preserve the existing scale so that adding source data does not change the model's input transform.

# Changes

- Preserve the first choice parameter's `log_scale` when merging two choice parameters, following the existing attribute precedence rule.
- Preserve the choice parameter's `log_scale` when merging it with a fixed parameter, in either argument order.
- Cover explicit logarithmic and linear scales with values for which automatic inference would choose the opposite scale.

# Examples

Input:
```python
choice = ChoiceParameter(
    "x", ParameterType.INT, values=[1, 2],
    is_ordered=True, sort_values=True, log_scale=True,
)
fixed = FixedParameter("x", ParameterType.INT, value=3)
merged = merge_parameters(choice, fixed, reverse_param_config={})
```

## Before

```python
merged.values     # [1, 2, 3]
merged.log_scale  # False: inferred again from the merged values
```

## After

```python
merged.values     # [1, 2, 3]
merged.log_scale  # True: preserved from the choice parameter
```

Likewise, merging an explicitly linear choice `[1, 2]` with fixed value `4` now preserves `False` instead of inferring `True` from `[1, 2, 4]`.

Differential Revision: D126675426
---
 .../transfer_learning/tests/test_utils.py     | 30 +++++++++++++++++++
 ax/adapter/transfer_learning/utils.py         |  2 ++
 2 files changed, 32 insertions(+)

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
             ),