Flow Annealed Importance Sampling Boostrap (FAB) implementation in JAX

See corresponding paper here.
See the pytorch implementation here.
The SNR analysis of the FAB gradient estimator is here.
I am currently working on a new version improved version of the jax implementation of FAB.

Examples

Double Well Boltzmann distribution samples vs contours (see notebook)

Gaussian Mixture Model samples vs contours (see notebook)

Name		Name	Last commit message	Last commit date
Latest commit History 145 Commits
examples_fabjax		examples_fabjax
fabjax		fabjax
flows_for_atomic_solids		flows_for_atomic_solids
.gitignore		.gitignore
MANIFEST.in		MANIFEST.in
README.md		README.md
install_req_tpu.sh		install_req_tpu.sh
jax_nan_playground.py		jax_nan_playground.py
requirements.txt		requirements.txt
setup.py		setup.py
vmap_playground.py		vmap_playground.py