Generative_models-flax Implementations of generative models using jax and flax. TODO VAEs GANs Diffusion models Train and generation on MNIST data Diffusion Process Visualization