Added JAX GAN demo in Colab

[susnato]

Description

Added JAX GAN demo (Google Colab Notebook Included)
Issue

#675

Checklist:

[✅] Performed a self-review of the code
[✅] Tested on Google Colab.

Potential problems/Important remarks

As described in this issue #675 I implemented GANs using JAX [Tensorflow was only used for Datasets purpose]with the help of https://github.com/valentingol/GANJax, then I converted that into a .py file. The Public Colab Notebook link: https://colab.research.google.com/drive/1YS-G6DML1xq3AhlNpKmh18Dc4tkn9T4U?usp=sharing

[susnato]

Can you please review my code? @mjsML @murphyk

[murphyk]

Some comments on your code.

• Since you want to commit a notebook, please make a PR at https://github.com/probml/probml-notebooks. (The pyprobml/scripts directory is reserved for smaller scripts that can be run locally on a laptop without a GPU, or for library code.)
• Put your name and contact info at top of the notebook. (It is in the .py file not by the .ipynb file).
• When you borrow code from someone else (eg GANjax repo), please be sure to mention this at the top of the colab, and specify what you have changed. (You do give proper credit, but it should appear at the top.)
• There is no need to include things like the JAX logo, or to repeat pedagogical material (eg picture of a GAN architeecture) that is already in the book. Similarly your summary 'Types Of GANs` is obviously very incomplete, and is not needed, since the book covers this already.
• You should only install what is necessary for colab, and avoid installing things that are already there.
  For example, this is overkill

-!pip install tensorflow -f https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow_cpu-2.8.0-cp39-cp39-manylinux2010_x86_64.whl -q
!pip install --upgrade pip -q
# Installs the wheel compatible with CUDA 11 and cuDNN 8.2 or newer.
!pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html  -q # Note: wheels only available on linux.
!pip install --upgrade pip -q

# Installs the wheel compatible with Cuda >= 11.4 and cudnn >= 8.2
!pip install jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_releases.html -q

# Installs the wheel compatible with Cuda >= 11.1 and cudnn >= 8.0.5
!pip install jax[cuda11_cudnn805] -f https://storage.googleapis.com/jax-releases/jax_releases.html -q
!sudo ln -s /path/to/cuda /usr/local/cuda-X.X
!export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/local/cuda/lib64
!pip install dm-haiku>=0.01 matplotlib>=3.5.0 optax>=0.1.0 -q

I found that this suffices

!pip install dm-haiku optax -q

• In the cell "Let's see some images from the dataset to make sure that they are ok.", the images look over-saturated (and the layout has an annoying blank space).

- You write "First let's try to train a model (you need to run for a long time 15 epochs at least) in order to get good looking outputs. Since you can't train your model in Colab for too long we are omitting that task but feel free to give it a try in your local machine". Most people will not have a GPU on their local machine. It would be great if you could make a detailed 'recipe' people could follow to train this model in GCP, and then just use Colab for visualizing it. - You write "I was facing some problems related to pickle so I had to rewrite load_jax_model.". What problems? - I was able to run the pre-trained model and get some samples, which is great. Maybe put this section before the section on training, so everyone can get some initial results quickly. - Please add section headers. -

[murphyk]

See also #691

[susnato]

Hi, @murphyk I made changes to the notebook and created another PR, here .So, I am closing this PR.