tpu_setup.sh

#!/bin/bash
# wget https://gist.githubusercontent.com/shawwn/88f64f7294c5a2e5e009d277a429ff2e/raw/tpu_setup.sh
# bash tpu_setup.sh
set -x

pip3 install --upgrade pip
# upgrade to nightly jax.
pip3 install --force-reinstall --pre -U -f https://storage.googleapis.com/jax-releases/libtpu_releases.html 'jax[tpu]' 'jaxlib'
pip3 install rich

if [ ! -f ~/.tmux.conf ]
then
cat >~/.tmux.conf <<'EOF'
set -g default-terminal "screen-256color"
set-option -g history-limit 300000
EOF
fi

if [ ! -x ~/.local/bin/userpath ]
then
  pip3 install userpath
  ~/.local/bin/userpath append ~/.local/bin
fi

if [ ! -d ~/scrap ]
then
  git clone https://github.com/shawwn/scrap ~/scrap
  ~/.local/bin/userpath prepend ~/scrap
  ~/.local/bin/userpath prepend ~/scrap/.bin
fi

sudo apt-get update
sudo apt-get install cmake rlwrap -y

mkdir -p ~/ml/data/tok16
if [ ! -d ~/ml/tlarc ]
then
  git clone https://github.com/tensorfork/tlarc ~/ml/tlarc
  cd ~/ml/tlarc
  bash bin/setup.sh
  bin/racket/bin/raco pkg install --auto unix-socket-lib
  cd ~/ml
fi


if [ ! -d ~/ml/diffusion ]
then
  pip3 install tensorflow-datasets tqdm fire
  git clone https://github.com/shawwn/diffusion ~/ml/diffusion
fi

if [ ! -d ~/ml/jax-exp ]
then
  pip3 install tabulate einops
  git clone https://github.com/shawwn/jax-exp ~/ml/jax-exp
fi

set +ex

cat <<'EOF'
# Now run:

exec $SHELL

# To train a diffusion model on a TPU:

cd ~/ml/diffusion
tmux new -s diffusion
python3 scripts/run_cifar.py train --exp_name run1
EOF