AutoEncoder models in Flax

This repository contains my implementation of various autoencoder models in Flax:

• The first notebook implements a vanilla autoencoder and a Variational AutoEncoder (VAE) on MNIST, with a comparison of the latent space representation in 2D.
  The models are also implemented in PyTorch for comparison
• The second notebook implements two methods to train the codebook of a Vector Quantised-Variational AutoEncoder on CIFAR-10, also provided in two standalone python files

References

1. VAE : Diederik P Kingma, Max Welling Auto-Encoding Variational Bayes, arXiv:1312.6114, December 2013
2. VQ VAE : Aaron van den Oord, Oriol Vinyals, Koray Kavukcuoglu Neural Discrete Representation Learning, arXiv:1711.00937, November 2017