Skip to content

Enforce equality constraints in random generator rejection sampling (#5182)#5182

Closed
sdaulton wants to merge 2 commits into
facebook:mainfrom
sdaulton:export-D100256485
Closed

Enforce equality constraints in random generator rejection sampling (#5182)#5182
sdaulton wants to merge 2 commits into
facebook:mainfrom
sdaulton:export-D100256485

Conversation

@sdaulton
Copy link
Copy Markdown
Contributor

@sdaulton sdaulton commented Apr 17, 2026

Summary:

Random generators (Sobol, etc.) were not respecting equality constraints
during candidate generation. Two fixes:

  1. When equality constraints are present, skip rejection sampling entirely
    and go straight to polytope sampling. Unconstrained random samples have
    probability zero of satisfying continuous equality constraints, so
    rejection sampling would always exhaust max_draws.

  2. Add equality_constraints parameter to rejection_sample and
    check_param_constraints so that post-rounding feasibility checks
    also validate equality constraints (important when the polytope sampler
    fallback uses rejection_sample for deduplication).

Reviewed By: esantorella

Differential Revision: D100256485

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

@meta-codesync meta-codesync Bot changed the title Enforce equality constraints in random generator rejection sampling Enforce equality constraints in random generator rejection sampling (#5182) Apr 17, 2026
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request Apr 17, 2026
…acebook#5182)

Summary:
Pull Request resolved: facebook#5182

Random generators (Sobol, etc.) were not respecting equality constraints
during candidate generation. Two fixes:

1. When equality constraints are present, skip rejection sampling entirely
   and go straight to polytope sampling. Unconstrained random samples have
   probability zero of satisfying continuous equality constraints, so
   rejection sampling would always exhaust max_draws.

2. Add `equality_constraints` parameter to `rejection_sample` and
   `check_param_constraints` so that post-rounding feasibility checks
   also validate equality constraints (important when the polytope sampler
   fallback uses rejection_sample for deduplication).

Differential Revision: D100256485
@sdaulton sdaulton force-pushed the export-D100256485 branch from 44be19a to 1da1880 Compare April 17, 2026 20:12
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request Apr 17, 2026
…acebook#5182)

Summary:
Pull Request resolved: facebook#5182

Random generators (Sobol, etc.) were not respecting equality constraints
during candidate generation. Two fixes:

1. When equality constraints are present, skip rejection sampling entirely
   and go straight to polytope sampling. Unconstrained random samples have
   probability zero of satisfying continuous equality constraints, so
   rejection sampling would always exhaust max_draws.

2. Add `equality_constraints` parameter to `rejection_sample` and
   `check_param_constraints` so that post-rounding feasibility checks
   also validate equality constraints (important when the polytope sampler
   fallback uses rejection_sample for deduplication).

Differential Revision: D100256485
@sdaulton sdaulton force-pushed the export-D100256485 branch from 1da1880 to e352e83 Compare April 17, 2026 20:41
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request Apr 17, 2026
…acebook#5182)

Summary:
Pull Request resolved: facebook#5182

Random generators (Sobol, etc.) were not respecting equality constraints
during candidate generation. Two fixes:

1. When equality constraints are present, skip rejection sampling entirely
   and go straight to polytope sampling. Unconstrained random samples have
   probability zero of satisfying continuous equality constraints, so
   rejection sampling would always exhaust max_draws.

2. Add `equality_constraints` parameter to `rejection_sample` and
   `check_param_constraints` so that post-rounding feasibility checks
   also validate equality constraints (important when the polytope sampler
   fallback uses rejection_sample for deduplication).

Differential Revision: D100256485
@sdaulton sdaulton force-pushed the export-D100256485 branch from e352e83 to cb20056 Compare April 17, 2026 22:34
@codecov-commenter
Copy link
Copy Markdown

codecov-commenter commented Apr 17, 2026

Codecov Report

❌ Patch coverage is 94.44444% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 96.38%. Comparing base (75c0ed4) to head (286fd3f).

Files with missing lines Patch % Lines
ax/generators/torch/botorch_modular/acquisition.py 92.68% 3 Missing ⚠️
ax/generators/utils.py 77.77% 2 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main    #5182      +/-   ##
==========================================
- Coverage   96.38%   96.38%   -0.01%     
==========================================
  Files         617      617              
  Lines       69494    69579      +85     
==========================================
+ Hits        66985    67065      +80     
- Misses       2509     2514       +5     

☔ 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#5182)

Summary:
Pull Request resolved: facebook#5182

Random generators (Sobol, etc.) were not respecting equality constraints
during candidate generation. Two fixes:

1. When equality constraints are present, skip rejection sampling entirely
   and go straight to polytope sampling. Unconstrained random samples have
   probability zero of satisfying continuous equality constraints, so
   rejection sampling would always exhaust max_draws.

2. Add `equality_constraints` parameter to `rejection_sample` and
   `check_param_constraints` so that post-rounding feasibility checks
   also validate equality constraints (important when the polytope sampler
   fallback uses rejection_sample for deduplication).

Differential Revision: D100256485
@sdaulton sdaulton force-pushed the export-D100256485 branch 2 times, most recently from 8679ff5 to 15f9dd9 Compare April 21, 2026 21:29
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request Apr 21, 2026
…acebook#5182)

Summary:

Random generators (Sobol, etc.) were not respecting equality constraints
during candidate generation. Two fixes:

1. When equality constraints are present, skip rejection sampling entirely
   and go straight to polytope sampling. Unconstrained random samples have
   probability zero of satisfying continuous equality constraints, so
   rejection sampling would always exhaust max_draws.

2. Add `equality_constraints` parameter to `rejection_sample` and
   `check_param_constraints` so that post-rounding feasibility checks
   also validate equality constraints (important when the polytope sampler
   fallback uses rejection_sample for deduplication).

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

Summary:
Pull Request resolved: facebook#5182

Random generators (Sobol, etc.) were not respecting equality constraints
during candidate generation. Two fixes:

1. When equality constraints are present, skip rejection sampling entirely
   and go straight to polytope sampling. Unconstrained random samples have
   probability zero of satisfying continuous equality constraints, so
   rejection sampling would always exhaust max_draws.

2. Add `equality_constraints` parameter to `rejection_sample` and
   `check_param_constraints` so that post-rounding feasibility checks
   also validate equality constraints (important when the polytope sampler
   fallback uses rejection_sample for deduplication).

Differential Revision: D100256485
@sdaulton sdaulton force-pushed the export-D100256485 branch from 15f9dd9 to a0c9d9c Compare April 21, 2026 21:34
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request May 1, 2026
…acebook#5182)

Summary:

Random generators (Sobol, etc.) were not respecting equality constraints
during candidate generation. Two fixes:

1. When equality constraints are present, skip rejection sampling entirely
   and go straight to polytope sampling. Unconstrained random samples have
   probability zero of satisfying continuous equality constraints, so
   rejection sampling would always exhaust max_draws.

2. Add `equality_constraints` parameter to `rejection_sample` and
   `check_param_constraints` so that post-rounding feasibility checks
   also validate equality constraints (important when the polytope sampler
   fallback uses rejection_sample for deduplication).

Differential Revision: D100256485
@sdaulton sdaulton force-pushed the export-D100256485 branch 2 times, most recently from 619938b to 5363507 Compare May 1, 2026 22:03
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request May 1, 2026
…acebook#5182)

Summary:

Random generators (Sobol, etc.) were not respecting equality constraints
during candidate generation. Two fixes:

1. When equality constraints are present, skip rejection sampling entirely
   and go straight to polytope sampling. Unconstrained random samples have
   probability zero of satisfying continuous equality constraints, so
   rejection sampling would always exhaust max_draws.

2. Add `equality_constraints` parameter to `rejection_sample` and
   `check_param_constraints` so that post-rounding feasibility checks
   also validate equality constraints (important when the polytope sampler
   fallback uses rejection_sample for deduplication).

Differential Revision: D100256485
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request May 5, 2026
…acebook#5182)

Summary:

Random generators (Sobol, etc.) were not respecting equality constraints
during candidate generation. Two fixes:

1. When equality constraints are present, skip rejection sampling entirely
   and go straight to polytope sampling. Unconstrained random samples have
   probability zero of satisfying continuous equality constraints, so
   rejection sampling would always exhaust max_draws.

2. Add `equality_constraints` parameter to `rejection_sample` and
   `check_param_constraints` so that post-rounding feasibility checks
   also validate equality constraints (important when the polytope sampler
   fallback uses rejection_sample for deduplication).

Reviewed By: esantorella

Differential Revision: D100256485
@sdaulton sdaulton force-pushed the export-D100256485 branch from 5363507 to 5840025 Compare May 5, 2026 16:18
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request May 5, 2026
…acebook#5182)

Summary:

Random generators (Sobol, etc.) were not respecting equality constraints
during candidate generation. Two fixes:

1. When equality constraints are present, skip rejection sampling entirely
   and go straight to polytope sampling. Unconstrained random samples have
   probability zero of satisfying continuous equality constraints, so
   rejection sampling would always exhaust max_draws.

2. Add `equality_constraints` parameter to `rejection_sample` and
   `check_param_constraints` so that post-rounding feasibility checks
   also validate equality constraints (important when the polytope sampler
   fallback uses rejection_sample for deduplication).

Reviewed By: esantorella

Differential Revision: D100256485
@sdaulton sdaulton force-pushed the export-D100256485 branch from 5840025 to 8e3d51d Compare May 5, 2026 16:20
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request May 5, 2026
…acebook#5182)

Summary:

Random generators (Sobol, etc.) were not respecting equality constraints
during candidate generation. Two fixes:

1. When equality constraints are present, skip rejection sampling entirely
   and go straight to polytope sampling. Unconstrained random samples have
   probability zero of satisfying continuous equality constraints, so
   rejection sampling would always exhaust max_draws.

2. Add `equality_constraints` parameter to `rejection_sample` and
   `check_param_constraints` so that post-rounding feasibility checks
   also validate equality constraints (important when the polytope sampler
   fallback uses rejection_sample for deduplication).

Reviewed By: esantorella

Differential Revision: D100256485
@sdaulton sdaulton force-pushed the export-D100256485 branch 2 times, most recently from d0c7002 to 020a277 Compare May 7, 2026 15:30
sdaulton added a commit to sdaulton/Ax-1 that referenced this pull request May 7, 2026
…acebook#5182)

Summary:

Random generators (Sobol, etc.) were not respecting equality constraints
during candidate generation. Two fixes:

1. When equality constraints are present, skip rejection sampling entirely
   and go straight to polytope sampling. Unconstrained random samples have
   probability zero of satisfying continuous equality constraints, so
   rejection sampling would always exhaust max_draws.

2. Add `equality_constraints` parameter to `rejection_sample` and
   `check_param_constraints` so that post-rounding feasibility checks
   also validate equality constraints (important when the polytope sampler
   fallback uses rejection_sample for deduplication).

Reviewed By: esantorella

Differential Revision: D100256485
sdaulton added 2 commits May 8, 2026 09:32
Summary:

Extend BONSAI's irrelevance pruning to handle both equality and inequality
constraints via a pin-and-project approach. Previously, BONSAI simply
discarded pruned candidates that violated constraints. This was overly
conservative (inequality) or completely broken (equality, where almost all
single-dimension prunes violate the constraint).

The new approach:
1. Set x_j = target[j] (unchanged)
2. Project the other dimensions onto the feasible set via SLSQP, keeping
   x_j pinned (and all previously pruned dims pinned)
3. Filter any candidates that remain infeasible after projection

This is strictly better than discarding: it recovers feasibility when
possible by adjusting other dimensions, while infeasible pins (where no
adjustment can satisfy the constraints) are still caught.

Key implementation details:
- `_project_and_filter_pruned_candidates`: new function that uses
  `project_to_feasible_space_via_slsqp` with `fixed_features` to pin
  the pruned dim and all previously pruned dims.
- Optimization: skip projection for dims not in any constraint's index
  set (pruning them can't violate anything).
- Handles 2D inter-point constraint indices correctly.
- `_prune_irrelevant_parameters` now accepts `bounds` parameter.

Reviewed By: esantorella

Differential Revision: D100256483
…acebook#5182)

Summary:

Random generators (Sobol, etc.) were not respecting equality constraints
during candidate generation. Two fixes:

1. When equality constraints are present, skip rejection sampling entirely
   and go straight to polytope sampling. Unconstrained random samples have
   probability zero of satisfying continuous equality constraints, so
   rejection sampling would always exhaust max_draws.

2. Add `equality_constraints` parameter to `rejection_sample` and
   `check_param_constraints` so that post-rounding feasibility checks
   also validate equality constraints (important when the polytope sampler
   fallback uses rejection_sample for deduplication).

Reviewed By: esantorella

Differential Revision: D100256485
@meta-codesync
Copy link
Copy Markdown

meta-codesync Bot commented May 8, 2026

This pull request has been merged in 47defa1.

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