fixing tf 2.18 install

Author: djherbis

tpu/Dockerfile

@@ -6,6 +6,7 @@ FROM $BASE_IMAGE
 # See: https://docs.docker.com/engine/reference/builder/#understand-how-arg-and-from-interact
 ARG PYTHON_WHEEL_VERSION
 ARG PYTHON_VERSION_PATH
+ARG TENSORFLOW_VERSION
 ARG TORCH_LINUX_WHEEL_VERSION
 ARG TORCH_VERSION
 ARG TORCHVISION_VERSION
@@ -61,9 +62,9 @@ RUN envsubst < /kaggle_requirements.txt > /requirements.txt
 
 # Install uv and then install the requirements:
 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 && \
+RUN export PATH="${HOME}/.local/bin:${PATH}" && uv pip install --system -r /requirements.txt --prerelease=allow --force-reinstall && \
     /tmp/clean-layer.sh
-ENV PATH="${HOME}/.local/bin:${PATH}"
+ENV PATH="~/.local/bin:${PATH}"
 
 # Kaggle Model Hub patches:
 ADD patches/kaggle_module_resolver.py /usr/local/lib/${PYTHON_VERSION_PATH}/site-packages/tensorflow_hub/kaggle_module_resolver.py

tpu/config.txt

@@ -1,7 +1,7 @@
 BASE_IMAGE=python:3.10
 PYTHON_WHEEL_VERSION=cp310
 PYTHON_VERSION_PATH=python3.10
-TF_LINUX_WHEEL_VERSION=manylinux_2_17_x86_64.manylinux2014_x86_64
+TENSORFLOW_VERSION=2.18.0
 # 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.5.0

tpu/requirements.txt

@@ -1,8 +1,8 @@
 # TPU Utils
 tpu-info
 # Tensorflow packages
-tensorflow-tpu>=2.18.0
--f https://storage.googleapis.com/libtpu-tf-releases/index.html
+tensorflow-tpu>=${TENSORFLOW_VERSION}
+--find-links https://storage.googleapis.com/libtpu-tf-releases/index.html
 tensorflow_hub
 tensorflow-io
 tensorflow-probability
@@ -13,6 +13,7 @@ torchaudio==${TORCHAUDIO_VERSION}
 torchvision==${TORCHVISION_VERSION}
 # Jax packages
 jax[tpu]>=0.4.34
+--find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
 distrax
 flax
 git+https://github.com/deepmind/dm-haiku