[Generative Adversarial Networks]

[nipunbatra]

See: https://code-first-ml.github.io/book2/notebooks/gans/2021-05-31-GAN.html

Rewrite in JAX/Flax