Skip to content

Support equality constraints in JSON and SQA storage (#5181)#5181

Open
sdaulton wants to merge 8 commits intofacebook:mainfrom
sdaulton:export-D100256489
Open

Support equality constraints in JSON and SQA storage (#5181)#5181
sdaulton wants to merge 8 commits intofacebook:mainfrom
sdaulton:export-D100256489

Conversation

@sdaulton
Copy link
Copy Markdown
Contributor

@sdaulton sdaulton commented Apr 17, 2026 •
edited by meta-codesync bot
Loading

Summary:

Update the JSON and SQA storage layers to correctly serialize and
deserialize equality constraints.

  • JSON encoder: use "equality" key with == for equality constraints,
    "inequality" key with <= for inequality constraints.
  • JSON decoder: detect "equality" key and construct
    ParameterConstraint(equality=...). Backward compatible — old format
    with "constraint_dict" + "bound" or "inequality" key still works.
  • SQA encoder: use ParameterConstraintType.EQUALITY for equality
    constraints, ParameterConstraintType.LINEAR for inequality.
  • SQA decoder: check type enum and construct with appropriate kwarg.
    No DB migration needed (new enum value, same column type).
  • Add EQUALITY = 4 to ParameterConstraintType enum.
  • Add __eq__ override to ParameterConstraint that compares semantic
    fields (constraint_dict, bound, is_equality) rather than raw
    constraint strings, ensuring round-trip equality.
  • Add equality_constraints to SobolGenerator.gen() and
    InSampleUniformGenerator.gen() signatures (they override base).

Differential Revision: D100256489

@meta-cla meta-cla bot added the CLA Signed Do not delete this pull request or issue due to inactivity. label Apr 17, 2026
@meta-codesync
Copy link
Copy Markdown

meta-codesync bot commented Apr 17, 2026

@sdaulton has exported this pull request. If you are a Meta employee, you can view the originating Diff in D100256489.

sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request Apr 17, 2026
@sdaulton @meta-codesync
Support equality constraints in JSON and SQA storage (facebook#5181)
39ccefb 
Summary:
Pull Request resolved: facebook#5181

Update the JSON and SQA storage layers to correctly serialize and
deserialize equality constraints.

- JSON encoder: use `"equality"` key with `==` for equality constraints,
  `"inequality"` key with `<=` for inequality constraints.
- JSON decoder: detect `"equality"` key and construct
  `ParameterConstraint(equality=...)`. Backward compatible — old format
  with `"constraint_dict"` + `"bound"` or `"inequality"` key still works.
- SQA encoder: use `ParameterConstraintType.EQUALITY` for equality
  constraints, `ParameterConstraintType.LINEAR` for inequality.
- SQA decoder: check type enum and construct with appropriate kwarg.
  No DB migration needed (new enum value, same column type).
- Add `EQUALITY = 4` to `ParameterConstraintType` enum.
- Add `__eq__` override to `ParameterConstraint` that compares semantic
  fields (`constraint_dict`, `bound`, `is_equality`) rather than raw
  constraint strings, ensuring round-trip equality.
- Add `equality_constraints` to `SobolGenerator.gen()` and
  `InSampleUniformGenerator.gen()` signatures (they override base).

Differential Revision: D100256489
@meta-codesync meta-codesync bot changed the title Support equality constraints in JSON and SQA storage Support equality constraints in JSON and SQA storage (#5181) Apr 17, 2026
@sdaulton sdaulton force-pushed the export-D100256489 branch from a7973ec to 39ccefb Compare April 17, 2026 20:11
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request Apr 17, 2026
@sdaulton @meta-codesync
Support equality constraints in JSON and SQA storage (facebook#5181)
7250a1e 
Summary:
Pull Request resolved: facebook#5181

Update the JSON and SQA storage layers to correctly serialize and
deserialize equality constraints.

- JSON encoder: use `"equality"` key with `==` for equality constraints,
  `"inequality"` key with `<=` for inequality constraints.
- JSON decoder: detect `"equality"` key and construct
  `ParameterConstraint(equality=...)`. Backward compatible — old format
  with `"constraint_dict"` + `"bound"` or `"inequality"` key still works.
- SQA encoder: use `ParameterConstraintType.EQUALITY` for equality
  constraints, `ParameterConstraintType.LINEAR` for inequality.
- SQA decoder: check type enum and construct with appropriate kwarg.
  No DB migration needed (new enum value, same column type).
- Add `EQUALITY = 4` to `ParameterConstraintType` enum.
- Add `__eq__` override to `ParameterConstraint` that compares semantic
  fields (`constraint_dict`, `bound`, `is_equality`) rather than raw
  constraint strings, ensuring round-trip equality.
- Add `equality_constraints` to `SobolGenerator.gen()` and
  `InSampleUniformGenerator.gen()` signatures (they override base).

Differential Revision: D100256489
@sdaulton sdaulton force-pushed the export-D100256489 branch from 39ccefb to 7250a1e Compare April 17, 2026 20:39
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request Apr 17, 2026
@sdaulton @meta-codesync
Support equality constraints in JSON and SQA storage (facebook#5181)
7272828 
Summary:
Pull Request resolved: facebook#5181

Update the JSON and SQA storage layers to correctly serialize and
deserialize equality constraints.

- JSON encoder: use `"equality"` key with `==` for equality constraints,
  `"inequality"` key with `<=` for inequality constraints.
- JSON decoder: detect `"equality"` key and construct
  `ParameterConstraint(equality=...)`. Backward compatible — old format
  with `"constraint_dict"` + `"bound"` or `"inequality"` key still works.
- SQA encoder: use `ParameterConstraintType.EQUALITY` for equality
  constraints, `ParameterConstraintType.LINEAR` for inequality.
- SQA decoder: check type enum and construct with appropriate kwarg.
  No DB migration needed (new enum value, same column type).
- Add `EQUALITY = 4` to `ParameterConstraintType` enum.
- Add `__eq__` override to `ParameterConstraint` that compares semantic
  fields (`constraint_dict`, `bound`, `is_equality`) rather than raw
  constraint strings, ensuring round-trip equality.
- Add `equality_constraints` to `SobolGenerator.gen()` and
  `InSampleUniformGenerator.gen()` signatures (they override base).

Differential Revision: D100256489
@sdaulton sdaulton force-pushed the export-D100256489 branch from 7250a1e to 7272828 Compare April 17, 2026 22:40
@codecov-commenter
Copy link
Copy Markdown

codecov-commenter commented Apr 17, 2026 •
edited
Loading

Codecov Report

❌ Patch coverage is 98.95397% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 96.42%. Comparing base (104693c) to head (b446a8b).

Files with missing lines Patch % Lines
ax/utils/common/sympy.py 76.92% 3 Missing ⚠️
ax/core/parameter_constraint.py 96.00% 1 Missing ⚠️
ax/storage/json_store/decoder.py 75.00% 1 Missing ⚠️
Additional details and impacted files 
@@            Coverage Diff             @@
##             main    #5181      +/-   ##
==========================================
+ Coverage   96.41%   96.42%   +0.01%     
==========================================
  Files         618      619       +1     
  Lines       68882    69310     +428     
==========================================
+ Hits        66410    66835     +425     
- Misses       2472     2475       +3

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

sdaulton added 8 commits April 20, 2026 09:40
@sdaulton @meta-codesync
Support equality parameter constraints in ParameterConstraint
0fa7084 
Summary:
Add support for linear equality constraints (`w^T x == b`) alongside
existing inequality constraints (`w^T x <= b`) in Ax's
`ParameterConstraint` class. This is the first diff in a stack that
threads equality constraints all the way down to BoTorch's
`optimize_acqf`.

Changes:
- Add `extract_coefficient_dict_from_equality` to `ax/utils/common/sympy.py`
  for parsing `"expr == bound"` strings (SymPy can't parse `==` directly
  since Python evaluates it as a boolean).
- Extend `ParameterConstraint.__init__` to accept `equality=` kwarg
  alongside existing `inequality=` kwarg. Exactly one must be provided.
- Add `is_equality` property.
- Update `check()` to use `|w^T x - b| <= tol` for equality constraints.
- Update `__repr__`, `clone()`, `clone_with_transformed_parameters()`.
- Add comprehensive tests for equality constraints.

Differential Revision: D100256486
@sdaulton @meta-codesync
Support equality constraints in service layer string parsing
ccb98b7 
Summary:
Add support for parsing equality constraint strings (e.g. `"x1 + x2 == 3"`)
in `constraint_from_str`. This extends the existing `<=`/`>=` parsing to also
accept `==` as a comparison operator.

- Add `_process_equality_constraint` function (analogous to `_process_linear_constraint`)
  that constructs `ParameterConstraint(equality=...)`.
- Detect `==` in `constraint_from_str` and route to the new function.
- Reject equality order constraints (`"x1 == x2"`) with a clear error message.
- Update `INVALID_CONSTRAINT_ERROR_MSG` to document `==` support.

Differential Revision: D100256487
@sdaulton @meta-codesync
Support equality constraints in SearchSpace membership checks
450d708 
Summary:
Update `SearchSpace.check_membership_df` and `compute_chebyshev_center` to
handle equality constraints alongside existing inequality constraints.

- `check_membership_df`: branch on `constraint.is_equality` — use
  `|weighted_sum - bound| <= tol` for equality, `weighted_sum <= bound + tol`
  for inequality.
- `compute_chebyshev_center`: separate equality and inequality constraints
  into `A_eq/b_eq` and `A_ub/b_ub` for `scipy.optimize.linprog`. Equality
  constraints don't get the `r * ||a_i||` augmentation since the inscribed
  ball center must lie on the hyperplane.
- `check_membership` already works via the updated `constraint.check()`.

Differential Revision: D100256478
@sdaulton @meta-codesync
Thread equality constraints through adapter layer and TorchOptConfig
545301b 
Summary:
Add equality constraint extraction and propagation through the adapter layer
to TorchOptConfig, enabling downstream generators to receive equality
constraints.

- Add `extract_equality_constraints` in `adapter_utils.py` (filters for
  `is_equality=True` constraints, returns `(A, b)` matrices).
- Update `extract_parameter_constraints` to filter out equality constraints.
- Add `equality_constraints` parameter to `validate_and_apply_final_transform`
  (now returns a 7-tuple).
- Add `equality_constraints: tuple[Tensor, Tensor] | None` field to
  `TorchOptConfig`.
- Update `TorchAdapter._get_transformed_model_gen_args` to extract equality
  constraints and pass them through to `TorchOptConfig`.

Differential Revision: D100256480
@sdaulton @meta-codesync
Wire equality constraints through generator and acquisition to BoTorch
8380a35 
Summary:
Thread equality constraints from TorchOptConfig through to BoTorch's
optimize_acqf and related optimizers. This is the key diff that connects
Ax's equality constraint representation to BoTorch's SLSQP-based optimizer.

- Add `_to_equality_constraints` in `torch/utils.py` — converts (A, b) tensor
  format to BoTorch's `(indices, coefficients, rhs)` format. No sign negation
  needed (equality is symmetric).
- Update `BoTorchGenerator.gen()` to pass equality constraints to
  `acqf.optimize()`.
- Add `equality_constraints` parameter to `Acquisition.optimize()` and forward
  to `optimize_acqf`, `optimize_acqf_mixed`, `optimize_acqf_mixed_alternating`.
- Raise `ValueError` for discrete optimizers and NSGA-II (unsupported).
- Update `validate_candidates` to check equality constraints.
- Update `_prune_irrelevant_parameters` and `_remove_infeasible_candidates`
  to handle equality constraints.
- Update `Surrogate.best_point` to pass equality constraints.

Differential Revision: D100256482
@sdaulton @meta-codesync
Support equality constraints in random generator polytope sampling
504401a 
Summary:
Thread equality constraints from the search space through the random adapter
to the `HitAndRunPolytopeSampler`. When rejection sampling falls back to
polytope sampling, equality constraints from both `fixed_features` and
parameter constraints are combined into a single `(C, c)` matrix.

- Update `RandomAdapter._gen` to extract equality constraints via
  `extract_equality_constraints`.
- Add `equality_constraints` parameter to `RandomGenerator.gen()`.
- Add `_combine_equality_constraints` method that merges fixed-feature-based
  and parameter-based equality constraints for the polytope sampler.

Differential Revision: D100256488
@sdaulton @meta-codesync
Handle equality constraints in UnitX transform
86cd224 
Summary:
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.

Differential Revision: D100256484
@sdaulton @meta-codesync
Support equality constraints in JSON and SQA storage (facebook#5181)
b446a8b 
Summary:
Pull Request resolved: facebook#5181

Update the JSON and SQA storage layers to correctly serialize and
deserialize equality constraints.

- JSON encoder: use `"equality"` key with `==` for equality constraints,
  `"inequality"` key with `<=` for inequality constraints.
- JSON decoder: detect `"equality"` key and construct
  `ParameterConstraint(equality=...)`. Backward compatible — old format
  with `"constraint_dict"` + `"bound"` or `"inequality"` key still works.
- SQA encoder: use `ParameterConstraintType.EQUALITY` for equality
  constraints, `ParameterConstraintType.LINEAR` for inequality.
- SQA decoder: check type enum and construct with appropriate kwarg.
  No DB migration needed (new enum value, same column type).
- Add `EQUALITY = 4` to `ParameterConstraintType` enum.
- Add `__eq__` override to `ParameterConstraint` that compares semantic
  fields (`constraint_dict`, `bound`, `is_equality`) rather than raw
  constraint strings, ensuring round-trip equality.
- Add `equality_constraints` to `SobolGenerator.gen()` and
  `InSampleUniformGenerator.gen()` signatures (they override base).

Differential Revision: D100256489
@sdaulton sdaulton force-pushed the export-D100256489 branch from 7272828 to b446a8b Compare April 20, 2026 16:46
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

CLA Signed Do not delete this pull request or issue due to inactivity. fb-exported meta-exported

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@sdaulton @codecov-commenter