flax

2f69b484 - Merge pull request #143 from jheek/fix-vae