bump tpu jax/torch versions, add tpu-info

Author: djherbis

tpu/Dockerfile

@@ -11,7 +11,6 @@ ARG TORCH_LINUX_WHEEL_VERSION
 ARG TORCH_VERSION
 ARG TENSORFLOW_VERSION
 ARG TF_LIBTPU_VERSION
-ARG JAX_VERSION
 ARG TORCHVISION_VERSION
 ARG TORCHAUDIO_VERSION
 
@@ -67,6 +66,7 @@ RUN envsubst < /kaggle_requirements.txt > /requirements.txt
 RUN curl -LsSf https://astral.sh/uv/install.sh | sh
 RUN export PATH="${HOME}/.local/bin:${PATH}" && uv pip install --system -r /requirements.txt --prerelease=allow --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html && \
     /tmp/clean-layer.sh
+ENV PATH="${HOME}/.local/bin:${PATH}"
 
 # Tensorflow libtpu:
 RUN curl --output /usr/local/lib/python3.10/site-packages/libtpu/libtpu.so https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/${TF_LIBTPU_VERSION}/libtpu.so

tpu/config.txt

@@ -6,12 +6,11 @@ PYTHON_VERSION_PATH=python3.10
 TENSORFLOW_VERSION=2.16.1
 TF_LIBTPU_VERSION=1.10.1
 TF_LINUX_WHEEL_VERSION=manylinux_2_17_x86_64.manylinux2014_x86_64
-JAX_VERSION=0.4.23
-# gsutil ls gs://pytorch-xla-releases/wheels/tpuvm/* | grep libtpu | grep -v -E ".*rc[0-9].*"
+# gsutil ls gs://pytorch-xla-releases/wheels/tpuvm/* | grep libtpu | grep torch_xla | grep -v -E ".*rc[0-9].*" | sed 's/.*torch_xla-\(.*\)+libtpu.*/\1/' | sort -rV
 # Supports nightly
-TORCH_VERSION=2.4.0
+TORCH_VERSION=2.5.0
 # https://github.com/pytorch/audio supports nightly
-TORCHAUDIO_VERSION=2.4.0
+TORCHAUDIO_VERSION=2.5.0
 # https://github.com/pytorch/vision supports nightly
-TORCHVISION_VERSION=0.19.0
+TORCHVISION_VERSION=0.20.0
 TORCH_LINUX_WHEEL_VERSION=manylinux_2_28_x86_64

tpu/requirements.txt

@@ -1,3 +1,5 @@
+# TPU Utils
+tpu-info
 # Tensorflow packages
 https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-${TENSORFLOW_VERSION}/tensorflow-${TENSORFLOW_VERSION}-${PYTHON_WHEEL_VERSION}-${PYTHON_WHEEL_VERSION}-${TF_LINUX_WHEEL_VERSION}.whl
 tensorflow_hub
@@ -9,7 +11,7 @@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-${TOR
 torchaudio==${TORCHAUDIO_VERSION}
 torchvision==${TORCHVISION_VERSION}
 # Jax packages
-jax[tpu]==${JAX_VERSION}
+jax[tpu]>=0.4.34
 distrax
 flax
 git+https://github.com/deepmind/dm-haiku