Skip to content

feat: Add AmortizedVIPosterior for amortized variational inference #1751

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
janfb wants to merge 21 commits into
base: main
Choose a base branch
from
Open

feat: Add AmortizedVIPosterior for amortized variational inference #1751

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
21 commits
Select commit Hold shift + click to select a range
f0e5e5f
feat: implement AmortizedVIPosterior and corresponding tests (WIP).
janfb Jan 13, 2026
3178b0b
feat: enhance AmortizedVIPosterior with vectorization and early stopp…
janfb Jan 13, 2026
e3dcb59
refactoring: input validation and error handling; enhance test for ac…
janfb Jan 13, 2026
4808920
feat: add FlowType enum for flow architectures in AmortizedVIPosterio…
janfb Jan 13, 2026
72e2bd5
feat: align amortized VI flow types and MAP behavior
janfb Jan 17, 2026
5bd06a1
refactoring: improve API consistency and input validation
janfb Feb 2, 2026
ea51371
feat(vi): add Zuko unconditional flow builder for VI
janfb Feb 2, 2026
7af45c0
refactor(vi): adapt DivergenceOptimizer for Zuko flow support
janfb Feb 2, 2026
4a1e186
refactor(vi): extend Zuko support in divergence optimizer subclasses
janfb Feb 2, 2026
4468d0d
fix(vi): address code review issues in VI posterior
janfb Feb 2, 2026
6982637
feat(vi): integrate Zuko unconditional flows into VIPosterior
janfb Feb 2, 2026
bc13be7
feat(vi): add train_amortized() method for amortized variational infe…
janfb Feb 2, 2026
5922057
fix(vi): fix serialization for Zuko flows and update tests
janfb Feb 2, 2026
fb837dd
refactor(vi): migrate amortized VI tests to unified VIPosterior
janfb Feb 2, 2026
e20a19e
refactor(vi): remove AmortizedVIPosterior from public API exports
janfb Feb 2, 2026
9875bfe
wip: fix zuko 1D issue, replace pyro flows with gaussian in 1D case
janfb Feb 4, 2026
14abaf2
device fixes and combining into one test file
janfb Feb 4, 2026
c53a83a
refactor tests: add fast amortized test; remove redundant.
janfb Feb 6, 2026
8b48387
clean up pyro; fix types; fix posterior params cls
janfb Feb 6, 2026
15206dd
Merge main into add-amortized-vip
janfb Feb 6, 2026
a040502
linting
janfb Feb 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 10 additions & 1 deletion sbi/inference/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
__all__ = _npe_family + _nre_family + _nle_family + _abc_family + ["FMPE", "NPSE"]

from sbi.inference.posteriors import (
AmortizedVIPosterior,
DirectPosterior,
EnsemblePosterior,
ImportanceSamplingPosterior,
Expand All @@ -53,4 +54,12 @@
)
from sbi.utils.simulation_utils import simulate_for_sbi

__all__ = ["FMPE", "MarginalTrainer", "NLE", "NPE", "NPSE", "NRE", "simulate_for_sbi"]
__all__ = [
"FMPE",
"MarginalTrainer",
"NLE",
"NPE",
"NPSE",
"NRE",
"simulate_for_sbi",
]
4 changes: 4 additions & 0 deletions sbi/inference/posteriors/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
# under the Apache License Version 2.0, see <https://www.apache.org/licenses/>

from sbi.inference.posteriors.amortized_vi_posterior import AmortizedVIPosterior
from sbi.inference.posteriors.direct_posterior import DirectPosterior
from sbi.inference.posteriors.ensemble_posterior import EnsemblePosterior
from sbi.inference.posteriors.importance_posterior import ImportanceSamplingPosterior
Expand All @@ -16,8 +17,11 @@
from sbi.inference.posteriors.rejection_posterior import RejectionPosterior
from sbi.inference.posteriors.vector_field_posterior import VectorFieldPosterior
from sbi.inference.posteriors.vi_posterior import VIPosterior
from sbi.neural_nets.factory import ZukoFlowType

__all__ = [
"AmortizedVIPosterior",
"ZukoFlowType",
"DirectPosterior",
"EnsemblePosterior",
"ImportanceSamplingPosterior",
Expand Down
Loading
Loading