Improve build_developer.sh

- Remove existing torch and torch_xla installations. When these packages
are installed from PyPI, it looks like `python setup.py` can't overwrite
them in all cases.

- Remove old wheels.

- Install pallas specific dependencies.

Author: tengyifei

scripts/build_developer.sh

@@ -7,6 +7,12 @@ set -x  # Display commands being run.
 cd "$(dirname "$(readlink -f "$0")")"
 cd ../../
 
+# First remove any left over old wheels
+# and old installation
+pip uninstall torch -y
+rm dist/*
+
+# Install pytorch
 python3 setup.py bdist_wheel
 python3 setup.py install
 cd ..
@@ -17,8 +23,10 @@ if [ -d "vision" ]; then
   python3 setup.py develop
 fi
 
+# Install torch_xla
 cd ..
 cd pytorch/xla
+pip uninstall torch_xla -y
 python3 setup.py develop
 
 # libtpu is needed to talk to the TPUs. If TPUs are not present,
@@ -27,6 +35,11 @@ pip install torch_xla[tpu] \
   -f https://storage.googleapis.com/libtpu-wheels/index.html \
   -f https://storage.googleapis.com/libtpu-releases/index.html
 
+# Install Pallas dependencies
+pip install torch_xla[pallas] \
+  -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html \
+  -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
+
 # Test that the library is installed correctly.
 python3 -c 'import torch_xla as xla; print(xla.device())'