Handle equality constraints in UnitX transform (#5179)

Summary:
Pull Request resolved: #5179

Update the UnitX transform to preserve equality constraint type when
transforming constraints to unit space. When constructing new constraints,
check `c.is_equality` and use `ParameterConstraint(equality=...)` for
equality constraints. The math is identical for both types (rescale
coefficients by `(u - l)` and adjust bound by `w * l`).

Other transforms (OneHot, IntToFloat, RemoveFixed, etc.) use `pc.clone()`
which automatically preserves `is_equality` — no changes needed.

Reviewed By: esantorella

Differential Revision: D100256484

fbshipit-source-id: 5195d354a269bfa790080a92ee47c32e405f1e4a

Author: sdaulton

ax/adapter/transforms/tests/test_unit_x_transform.py

@@ -133,6 +133,40 @@ def test_TransformSearchSpace(self) -> None:
             self.search_space_with_target.parameters["x"].target_value, 1.0
         )
 
+    def test_TransformSearchSpaceEqualityConstraints(self) -> None:
+        # Verify that equality constraints are preserved and correctly
+        # rescaled during UnitX transform_search_space.
+        ss = SearchSpace(
+            parameters=[
+                RangeParameter(
+                    "x", lower=1, upper=3, parameter_type=ParameterType.FLOAT
+                ),
+                RangeParameter(
+                    "y", lower=1, upper=2, parameter_type=ParameterType.FLOAT
+                ),
+            ],
+            parameter_constraints=[
+                ParameterConstraint(equality="-0.5*x + y == 0.5"),
+                ParameterConstraint(inequality="-0.5*x + y <= 0.5"),
+            ],
+        )
+        t = UnitX(search_space=ss)
+        ss = t.transform_search_space(ss)
+
+        # Both constraints have the same math; only the type differs.
+        eq_c = ss.parameter_constraints[0]
+        ineq_c = ss.parameter_constraints[1]
+
+        # Equality constraint preserved.
+        self.assertTrue(eq_c.is_equality)
+        self.assertEqual(eq_c.constraint_dict, {"x": -1.0, "y": 1.0})
+        self.assertEqual(eq_c.bound, 0.0)
+
+        # Inequality constraint preserved.
+        self.assertFalse(ineq_c.is_equality)
+        self.assertEqual(ineq_c.constraint_dict, {"x": -1.0, "y": 1.0})
+        self.assertEqual(ineq_c.bound, 0.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

ax/adapter/transforms/unit_x.py

@@ -103,11 +103,14 @@ def transform_search_space(self, search_space: SearchSpace) -> SearchSpace:
             expr = " + ".join(
                 f"{coeff} * {param}" for param, coeff in constraint_dict.items()
             )
-            new_constraints.append(
-                ParameterConstraint(
-                    inequality=f"{expr} <= {bound}",
+            if c.is_equality:
+                new_constraints.append(
+                    ParameterConstraint(equality=f"{expr} == {bound}")
+                )
+            else:
+                new_constraints.append(
+                    ParameterConstraint(inequality=f"{expr} <= {bound}")
                 )
-            )
         search_space.set_parameter_constraints(new_constraints)
         return search_space
 

ax/generators/tests/test_sobol.py

@@ -305,20 +305,20 @@ def test_SobolGeneratorFallbackToPolytopeSamplerWithEqAndFixedFeatures(
                 )
             except Exception:
                 pass
-            if MockSampler.called:
-                eq_arg = MockSampler.call_args.kwargs["equality_constraints"]
-                self.assertIsNotNone(eq_arg)
-                C, c = eq_arg
-                # 1 from fixed features + 1 from parameter equality = 2 rows.
-                self.assertEqual(C.shape[0], 2)
-                self.assertEqual(C.shape[1], 11)
-                # Fixed feature: x[10] = 1 (last row from fixed, first in sorted).
-                self.assertEqual(C[0, 10].item(), 1.0)
-                self.assertAlmostEqual(c[0].item(), 1.0)
-                # Parameter equality: x[0] + x[1] = 0.3.
-                self.assertEqual(C[1, 0].item(), 1.0)
-                self.assertEqual(C[1, 1].item(), 1.0)
-                self.assertAlmostEqual(c[1].item(), 0.3)
+            self.assertTrue(MockSampler.called)
+            eq_arg = MockSampler.call_args.kwargs["equality_constraints"]
+            self.assertIsNotNone(eq_arg)
+            C, c = eq_arg
+            # 1 from fixed features + 1 from parameter equality = 2 rows.
+            self.assertEqual(C.shape[0], 2)
+            self.assertEqual(C.shape[1], 11)
+            # Fixed feature: x[10] = 1 (last row from fixed, first in sorted).
+            self.assertEqual(C[0, 10].item(), 1.0)
+            self.assertAlmostEqual(c[0].item(), 1.0)
+            # Parameter equality: x[0] + x[1] = 0.3.
+            self.assertEqual(C[1, 0].item(), 1.0)
+            self.assertEqual(C[1, 1].item(), 1.0)
+            self.assertAlmostEqual(c[1].item(), 0.3)
 
     def test_SobolGeneratorFallbackToPolytopeSamplerWithFixedParam(self) -> None:
         # Ten parameters with sum less than 1. In this example, the rejection