flax

6a355b22 - VAE with a discrete bottleneck with quantization done using the kmeans algorithm.