[Thanatos9404]

Motivation

Fixes #1687.

optimize_acqf(..., options={"sample_around_best": True}) fails for a HigherOrderGP with IndexError: tuple index out of range (botorch/optim/initializers.py:1249 on main). With a multi-output objective (e.g. HOGP with qNEHVI) it does not raise, but silently perturbs the wrong training points: in one check main used row [1] while the Pareto set is [4, 9].

Cause. sample_points_around_best averages the posterior mean over every dimension but the last two (while mean.ndim > 2: mean = mean.mean(dim=0)), assuming a batch_shape x n x m mean. A HigherOrderGP mean is batch_shape x n x output_shape, so the loop averages away the n baseline points.

Changes (both in sample_points_around_best):

1. Average only over len(acq_function.model.batch_shape) leading dimensions, per the Model.batch_shape contract. If batch_shape raises NotImplementedError (e.g. EnsembleModel) or is None (e.g. GenericDeterministicModel), the previous rule is used. The count is capped at mean.ndim - 2, so it never averages over more dimensions than before.
2. Treat the selection as multi-objective when f_pred is n x m (f_pred.ndim == 2 and f_pred.shape[-1] > 1) instead of comparing with mean.ndim. For n x m means the two checks are equivalent.

For every other model len(model.batch_shape) == mean.ndim - 2, so behaviour is unchanged. Comparing main and this branch side by side (SingleTaskGP plain/multi-output/batched, ModelListGP, MultiTaskGP, KroneckerMultiTaskGP, SAAS fully Bayesian, EnsembleModel, GenericDeterministicModel, MixedSingleTaskGP, multi-fidelity, variational, with transforms, scalarized and constrained multi-objective acquisitions) gave identical selected points; only HigherOrderGP changes.

Have you read the Contributing Guidelines on pull requests?

The change follows them: unit tests added, ufmt and flake8 clean, no public API or docs changes. The Meta CLA will be signed when the bot asks.

Test Plan

New tests in test/optim/test_initializers.py (TestSampleAroundBest):

• test_sample_points_around_best_higher_order_gp: HOGP with qEI (best point under the posterior mean), gen_batch_initial_conditions end to end as in the issue, and qNEHVI (exactly the Pareto-optimal training points).
• test_sample_points_around_best_batched_model: pins the existing behaviour for a batched SingleTaskGP and for MockModels whose batch_shape is larger than the mean, raises NotImplementedError, or is None.

pytest test/optim/test_initializers.py -k TestSampleAroundBest: on main 1 failed (the HOGP test, with the issue's IndexError), 2 passed; with this change 3 passed. test/optim/ and the HOGP model, posterior and multi-objective acquisition tests pass. ufmt check and flake8 are clean.

Related PRs

No docs changes. #3395 changes a different line of the same function and merges cleanly with this PR.

This change was written with Claude Code (Claude Opus 5.5) under my direction, checked by a separate AI review pass and tested only as described above, and I authorized submitting it.

[meta-cla]

Hi @Thanatos9404!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[meta-cla]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!