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)