Get tensorflow, jax, and pytorch working on TPU1VM

http://b/213335159

Author: djherbis

tpu/Dockerfile

@@ -10,4 +10,13 @@ COPY --from=libtpu /libtpu.so /lib
 
 COPY --from=tensorflow_whl /tmp/tensorflow_pkg/tensorflow*.whl /tmp/tensorflow_pkg/
 RUN pip install /tmp/tensorflow_pkg/tensorflow*.whl && \
-    rm -rf /tmp/tensorflow_pkg
+    rm -rf /tmp/tensorflow_pkg
+
+# https://cloud.google.com/tpu/docs/pytorch-xla-ug-tpu-vm#changing_pytorch_version
+RUN pip uninstall -y torch
+RUN pip uninstall -y torch_xla
+RUN pip install torch==1.10
+RUN pip install torch_xla[tpuvm] -f https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-1.10-cp37-cp37m-linux_x86_64.whl
+
+# https://cloud.google.com/tpu/docs/jax-quickstart-tpu-vm#install_jax_on_your_cloud_tpu_vm
+RUN pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

tpu/tensorflow.Dockerfile

@@ -17,7 +17,7 @@ RUN cd /usr/local/src && \
     pip install keras_applications --no-deps && \
     pip install keras_preprocessing --no-deps
 
-# Create a TensorFlow wheel for CPU
+# Create a TensorFlow wheel for TPU
 RUN cd /usr/local/src/tensorflow && \
     cat /dev/null | ./configure && \
     bazel build \
@@ -32,7 +32,22 @@ RUN cd /usr/local/src/tensorflow && \
 RUN cd /usr/local/src/tensorflow && \
     bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/tensorflow_pkg
 
-# TODO(b/152075195): Will likely need to install custom build for TFA & tensorflow-gcs-config
+# Build TensorFlow addons library against TensorFlow CPU.
+#RUN cd /usr/local/src/ && \
+#    git clone https://github.com/tensorflow/addons && \
+#    cd addons && \
+#    git checkout tags/v0.12.1 && \
+#    python ./configure.py && \
+#    bazel build --enable_runfiles build_pip_pkg && \
+#    bazel-bin/build_pip_pkg /tmp/tfa_cpu && \
+#    bazel clean
+
+# Build tensorflow_gcs_config library against TensorFlow CPU.
+#ADD tensorflow-gcs-config /usr/local/src/tensorflow_gcs_config/
+#RUN cd /usr/local/src/tensorflow_gcs_config && \
+#    apt-get install -y libcurl4-openssl-dev && \
+#    python setup.py bdist_wheel -d /tmp/tensorflow_gcs_config && \
+#    bazel clean
 
 # Use multi-stage builds to minimize image output size.
 FROM alpine:latest