Skip to content

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

Open
sdaulton wants to merge 3 commits intofacebook:mainfrom
sdaulton:export-D100256485
Open

Enforce equality constraints in random generator rejection sampling (#5182)#5182
sdaulton wants to merge 3 commits intofacebook: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).

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.04762% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 96.41%. Comparing base (8c8dc6c) to head (a0c9d9c).
⚠️ Report is 1 commits behind head on main.

Files with missing lines Patch % Lines
ax/generators/utils.py 77.77% 2 Missing ⚠️
ax/core/parameter_constraint.py 75.00% 1 Missing ⚠️
ax/generators/torch/botorch_modular/acquisition.py 97.14% 1 Missing ⚠️
ax/storage/json_store/decoder.py 75.00% 1 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main    #5182      +/-   ##
==========================================
- Coverage   96.42%   96.41%   -0.02%     
==========================================
  Files         619      619              
  Lines       69293    69371      +78     
==========================================
+ Hits        66818    66883      +65     
- Misses       2475     2488      +13     

☔ 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 from cb20056 to 8679ff5 Compare April 20, 2026 16:45
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
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.

Differential Revision: D100256483
@sdaulton sdaulton force-pushed the export-D100256485 branch 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
…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
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 meta-exported

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants