Skip to content

Support equality constraints in random generator polytope sampling (#5178)#5178

Closed
sdaulton wants to merge 6 commits intofacebook:mainfrom
sdaulton:export-D100256488
Closed

Support equality constraints in random generator polytope sampling (#5178)#5178
sdaulton wants to merge 6 commits intofacebook:mainfrom
sdaulton:export-D100256488

Conversation

@sdaulton
Copy link
Copy Markdown
Contributor

@sdaulton sdaulton commented Apr 17, 2026

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.

Reviewed By: bletham

Differential Revision: D100256488

@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 D100256488.

@meta-codesync meta-codesync Bot changed the title Support equality constraints in random generator polytope sampling Support equality constraints in random generator polytope sampling (#5178) Apr 17, 2026
@sdaulton sdaulton force-pushed the export-D100256488 branch from d481bea to 36f2035 Compare April 17, 2026 20:09
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request Apr 17, 2026
…acebook#5178)

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.

Reviewed By: bletham

Differential Revision: D100256488
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request Apr 17, 2026
…acebook#5178)

Summary:
Pull Request resolved: facebook#5178

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.

Reviewed By: bletham

Differential Revision: D100256488
@sdaulton sdaulton force-pushed the export-D100256488 branch 2 times, most recently from 31ccb4c to 5ffa36e Compare April 17, 2026 20:36
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request Apr 17, 2026
…acebook#5178)

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.

Reviewed By: bletham

Differential Revision: D100256488
@sdaulton sdaulton force-pushed the export-D100256488 branch from 5ffa36e to 64fd10e Compare April 17, 2026 20:40
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request Apr 17, 2026
…acebook#5178)

Summary:
Pull Request resolved: facebook#5178

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.

Reviewed By: bletham

Differential Revision: D100256488
@sdaulton sdaulton force-pushed the export-D100256488 branch from 64fd10e to f823d41 Compare April 17, 2026 22:31
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request Apr 17, 2026
…acebook#5178)

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.

Reviewed By: bletham

Differential Revision: D100256488
@sdaulton sdaulton force-pushed the export-D100256488 branch from f823d41 to baadb4f Compare April 17, 2026 22:34
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request Apr 17, 2026
…acebook#5178)

Summary:
Pull Request resolved: facebook#5178

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.

Reviewed By: bletham

Differential Revision: D100256488
@codecov-commenter
Copy link
Copy Markdown

codecov-commenter commented Apr 17, 2026

Codecov Report

❌ Patch coverage is 99.32584% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 96.42%. Comparing base (104693c) to head (317fe1e).
⚠️ Report is 7 commits behind head on main.

Files with missing lines Patch % Lines
ax/utils/common/sympy.py 76.92% 3 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main    #5178      +/-   ##
==========================================
+ Coverage   96.41%   96.42%   +0.01%     
==========================================
  Files         618      619       +1     
  Lines       68882    69279     +397     
==========================================
+ Hits        66410    66803     +393     
- Misses       2472     2476       +4     

☔ 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 a commit to sdaulton/Ax-1 that referenced this pull request Apr 20, 2026
…acebook#5178)

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.

Reviewed By: bletham

Differential Revision: D100256488
@sdaulton sdaulton force-pushed the export-D100256488 branch 2 times, most recently from ddf8e41 to 4035b66 Compare April 20, 2026 16:46
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request Apr 20, 2026
…acebook#5178)

Summary:
Pull Request resolved: facebook#5178

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.

Reviewed By: bletham

Differential Revision: D100256488
…ok#5173)

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.

Reviewed By: esantorella

Differential Revision: D100256486
…k#5174)

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.

Reviewed By: bletham

Differential Revision: D100256487
…ok#5175)

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()`.

Reviewed By: bletham

Differential Revision: D100256478
…acebook#5176)

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`.

Reviewed By: bletham

Differential Revision: D100256480
facebook#5177)

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.

Reviewed By: bletham

Differential Revision: D100256482
…acebook#5178)

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.

Reviewed By: bletham

Differential Revision: D100256488
@sdaulton sdaulton force-pushed the export-D100256488 branch from 4035b66 to 317fe1e Compare April 20, 2026 20:54
@meta-codesync
Copy link
Copy Markdown

meta-codesync Bot commented Apr 20, 2026

This pull request has been merged in 1c8c06e.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

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

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants