make pedagogical GAN example in 2d

[murphyk]

Implement a GAN from scratch (in JAX) to fit a dataset generated from a mixture of Gaussians in 2d.
You can start with this:
https://colab.research.google.com/github/probml/probml-notebooks/blob/master/notebooks/gan_mog_mode_hopping.ipynb
but your code should have these changes:

• it should be simpler (not use ODEs)
• use flax not haiku
• do not mix np.random and jax.random

[workatcam]

Hello Kevin, thank you for providing new tasks. I’d love to make pedagogical GAN from scratch using JAX, can I take over this issue if you don’t mind :)

[murphyk]

Hi @workatcam . Please go for it :)