forked from pytorch/botorch
-
Notifications
You must be signed in to change notification settings - Fork 1
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Extend Posterior API to support torch distributions & overhaul MCSamp…
…ler API (pytorch#1254) Summary: X-link: facebook/Ax#1254 X-link: facebookresearch/aepsych#193 Pull Request resolved: pytorch#1486 The main goal here is to broadly support non-Gaussian posteriors. - Adds a generic `TorchPosterior` which wraps a Torch `Distribution`. This defines a few properties that we commonly expect, and calls the `distribution` for the rest. - For a unified plotting API, this shifts away from mean & variance to a quantile function. Most torch distributions implement inverse CDF, which is used as quantile. For others, the user should implement it either at distribution or posterior level. - Hands off the burden of base sample handling from the posterior to the samplers. Using a dispatcher based `get_sampler` method, we can support SAA with mixed posteriors without having to shuffle base samples in a `PosteriorList`, as long as all base distributions have a corresponding sampler and support base samples. - Adds `ListSampler` for sampling from `PosteriorList`. - Adds `ForkedRNGSampler` and `StochasticSampler` for sampling from posteriors without base samples. - Adds `rsample_from_base_samples` for sampling with `base_samples` / with a `sampler`. - Absorbs `FullyBayesianPosteriorList` into `PosteriorList`. - For MC acqfs, introduces a `get_posterior_samples` for sampling from the posterior with base samples / a sampler. If a sampler was not specified, this constructs the appropriate sampler for the posterior using `get_sampler`, eliminating the need to construct a sampler in `__init__`, which we used to do under the assumption of Gaussian posteriors. TODOs: - Relax the Gaussian assumption in acquisition functions & utilities. Some of this might be addressed in a follow-up diff. - Updates to website / docs & tutorials to clear up some of the Gaussian assumption, introduce the new relaxed API. Likely a follow-up diff. - Some more listed in T134364907 - Test fixes and new units Other notables: - See D39760855 for usage of TorchDistribution in SkewGP. - TransformedPosterior could serve as the fallback option for derived posteriors. - MC samplers no longer support resample or collapse_batch_dims(=False). These can be handled by i) not using base samples, ii) just using torch.fork_rng and sampling without base samples from that. Samplers are only meant to support SAA. Introduces `ForkedRNGSampler` and `StochasticSampler` as convenience samplers for these use cases. - Introduced `batch_range_override` for the sampler to support edge cases where we may want to override `posterior.batch_range` (needed in `qMultiStepLookahead`) - Removes unused sampling utilities `construct_base_samples(_from_posterior)`, which assume Gaussian posterior. - Moves the main logic of `_set_sampler` method of CachedCholesky subclasses to a `_update_base_samples` method on samplers, and simplifies these classes a bit more. Reviewed By: Balandat Differential Revision: D39759489 fbshipit-source-id: f1a7d9b390e4012a35f13cfd04ee1114bc12e53e
- Loading branch information
1 parent
24ee799
commit 15d6cdc
Showing
99 changed files
with
2,735 additions
and
2,886 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.