Upgrade TPU Pytorch to 1.13

Author: djherbis

tpu/Dockerfile

@@ -4,6 +4,10 @@ FROM python:3.8
 # See: https://docs.docker.com/engine/reference/builder/#understand-how-arg-and-from-interact
 ARG TORCH_VERSION
 ARG TENSORFLOW_VERSION
+ARG JAX_VERSION
+ARG TORCHVISION_VERSION
+ARG TORCHTEXT_VERSION
+ARG TORCHAUDIO_VERSION
 
 ENV ISTPUVM=1
 
@@ -27,14 +31,14 @@ ENV JAX_LIBTPU=/lib/jax-libtpu.so
 
 # Install JAX & related packages
 # https://cloud.google.com/tpu/docs/jax-quickstart-tpu-vm#install_jax_on_your_cloud_tpu_vm
-RUN pip install "jax[tpu]==0.3.10" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html trax flax optax elegy git+https://github.com/deepmind/dm-haiku jraph distrax
+RUN pip install jax[tpu]==${JAX_VERSION} -f https://storage.googleapis.com/jax-releases/libtpu_releases.html trax flax optax elegy git+https://github.com/deepmind/dm-haiku jraph distrax
 
 RUN cp $PIP_LIBTPU $JAX_LIBTPU
 
 # Install Pytorch & related packages
 # https://cloud.google.com/tpu/docs/pytorch-xla-ug-tpu-vm#changing_pytorch_version
 # The URL doesn't include patch version. i.e. must use 1.11 instead of 1.11.0
-RUN pip install torch==${TORCH_VERSION} torch_xla[tpuvm] -f https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-${TORCH_VERSION%.*}-cp38-cp38-linux_x86_64.whl torchvision==0.12.0 torchtext==0.12.0 torchaudio==0.11.0
+RUN pip install torch==${TORCH_VERSION} torch_xla[tpuvm] -f https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-${TORCH_VERSION%.*}-cp38-cp38-linux_x86_64.whl torchvision==${TORCHVISION_VERSION} torchtext==${TORCHTEXT_VERSION} torchaudio==${TORCHAUDIO_VERSION}
 
 RUN cp $PIP_LIBTPU $PYTORCH_LIBTPU
 

tpu/config.txt

@@ -1,2 +1,10 @@
 TENSORFLOW_VERSION=2.9.1
-TORCH_VERSION=1.11.0
+JAX_VERSION=0.3.10
+# Supports nightly
+TORCH_VERSION=1.13.0
+# https://github.com/pytorch/audio supports nightly
+TORCHAUDIO_VERSION=0.13.0
+# https://github.com/pytorch/text supports main
+TORCHTEXT_VERSION=0.14.0
+# https://github.com/pytorch/vision supports nightly
+TORCHVISION_VERSION=0.14.0