A simple playground for testing Jax at Mila.

Install Python:

• module load python/3.10 (add to .profile)
• cd ~
• python -m venv python_environment
• source $HOME/python_environment/bin/activate (add to .profile)

Install cuDNN:

• cd ~
• tar --strip-components=1 -xf /network/scratch/b/bilaniuo/cuDNN/v8.6.0.163/cudnn-linux-x86_64-8.6.0.163_cuda11-archive.tar.xz -C ~/cudnn
• export CPATH="$HOME/cudnn/include:$CPATH (add to .profile)
• export LD_LIBRARY_PATH="$HOME/cudnn/lib:$LD_LIBRARY_PATH (add to .profile)

Install poetry:

• pip install poetry
• alias poetry=$HOME/python_environment/bin/poetry (add to .bashrc)

Install this repository:

• git clone https://github.com/NeilGirdhar/mila_jax.git
• cd mila_jax
• poetry install

Run the test:

• cd mila_jax
• poetry shell
• sbatch job.sh