Skip to content

Wire equality constraints through generator and acquisition to BoTorch (#5177)#5177

Closed
sdaulton wants to merge 5 commits intofacebook:mainfrom
sdaulton:export-D100256482
Closed

Wire equality constraints through generator and acquisition to BoTorch (#5177)#5177
sdaulton wants to merge 5 commits intofacebook:mainfrom
sdaulton:export-D100256482

Conversation

@sdaulton
Copy link
Copy Markdown
Contributor

@sdaulton sdaulton commented Apr 17, 2026

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

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

sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request Apr 17, 2026
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
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request Apr 17, 2026
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
@meta-codesync meta-codesync Bot changed the title Wire equality constraints through generator and acquisition to BoTorch Wire equality constraints through generator and acquisition to BoTorch (#5177) Apr 17, 2026
@sdaulton sdaulton force-pushed the export-D100256482 branch from 7164e29 to ce8c18d Compare April 17, 2026 20:09
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request Apr 17, 2026
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
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request Apr 17, 2026
facebook#5177)

Summary:
Pull Request resolved: facebook#5177

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
@sdaulton sdaulton force-pushed the export-D100256482 branch from ce8c18d to d63d175 Compare April 17, 2026 20:13
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request Apr 17, 2026
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
@sdaulton sdaulton force-pushed the export-D100256482 branch from d63d175 to 74cf46b Compare April 17, 2026 20:37
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request Apr 17, 2026
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
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request Apr 17, 2026
facebook#5177)

Summary:
Pull Request resolved: facebook#5177

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
@sdaulton sdaulton force-pushed the export-D100256482 branch from 74cf46b to 1dbb592 Compare April 17, 2026 20:40
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request Apr 17, 2026
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
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request Apr 20, 2026
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
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request Apr 20, 2026
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
@sdaulton sdaulton force-pushed the export-D100256482 branch from 1dbb592 to 52657f3 Compare April 20, 2026 16:43
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request Apr 20, 2026
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
@sdaulton sdaulton force-pushed the export-D100256482 branch from 52657f3 to cea1e4b Compare April 20, 2026 16:45
@codecov-commenter
Copy link
Copy Markdown

codecov-commenter commented Apr 20, 2026

Codecov Report

❌ Patch coverage is 99.12536% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 96.42%. Comparing base (104693c) to head (842866a).
⚠️ 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    #5177    +/-   ##
========================================
  Coverage   96.41%   96.42%            
========================================
  Files         618      619     +1     
  Lines       68882    69177   +295     
========================================
+ Hits        66410    66701   +291     
- 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.

…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
@sdaulton sdaulton force-pushed the export-D100256482 branch from cea1e4b to 842866a Compare April 20, 2026 20:53
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request Apr 20, 2026
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
@meta-codesync
Copy link
Copy Markdown

meta-codesync Bot commented Apr 20, 2026

This pull request has been merged in 58cff45.

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