Awesome

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) Double Well samples vs contours

Gaussian Mixture Model samples vs contours (see notebook)