TPU upgrades to tf and torch, clean up dead code (#1480)

Since the latest colab image is now python 3.11, we should try to bump
this as well.

Also removing comments and code that are out of date

Author: calderjo

tpu/Dockerfile

@@ -25,36 +25,13 @@ ADD patches/kaggle_session.py /root/.local/lib/${PYTHON_VERSION_PATH}/site-packa
 ADD patches/kaggle_web_client.py /root/.local/lib/${PYTHON_VERSION_PATH}/site-packages/kaggle_web_client.py
 ADD patches/kaggle_datasets.py /root/.local/lib/${PYTHON_VERSION_PATH}/site-packages/kaggle_datasets.py
 
-# Disable GCP integrations for now
-# ADD patches/kaggle_gcp.py /root/.local/lib/${PYTHON_VERSION_PATH}/site-packages/kaggle_gcp.py
-
-# Disable logging to file (why do we need this?)
-# ADD patches/log.py /root/.local/lib/${PYTHON_VERSION_PATH}/site-packages/log.py
-
-# sitecustomize adds significant latency to ipython kernel startup and should only be added if needed
-# ADD patches/sitecustomize.py /root/.local/lib/${PYTHON_VERSION_PATH}/site-packages/sitecustomize.py
-
 # Prereqs
 # This is needed for cv2 (opencv-python):
 # https://stackoverflow.com/questions/55313610/importerror-libgl-so-1-cannot-open-shared-object-file-no-such-file-or-directo
 RUN apt-get update && apt-get install ffmpeg libsm6 libxext6  -y
 
 # Install all the packages together for maximum compatibility.
-
-# Install Tensorflow.
-
-# 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
-# We need to keep the numpy version the same as the installed tf one but compatible with other installs.
-
-# Install JAX & related packages
-# https://cloud.google.com/tpu/docs/jax-quickstart-tpu-vm#install_jax_on_your_cloud_tpu_vm
-
-# Packages needed by the Notebook editor
-
 # Additional useful packages should be added in the requirements.txt
-
 # Bring in the requirements.txt and replace variables in it:
 RUN apt-get install -y gettext
 ADD tpu/requirements.txt /kaggle_requirements.txt

tpu/config.txt

@@ -4,9 +4,9 @@ PYTHON_VERSION_PATH=python3.10
 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
+TORCH_VERSION=2.6.0
 # https://github.com/pytorch/audio supports nightly
-TORCHAUDIO_VERSION=2.5.0
+TORCHAUDIO_VERSION=2.6.0
 # https://github.com/pytorch/vision supports nightly
-TORCHVISION_VERSION=0.20.0
+TORCHVISION_VERSION=0.21.0
 TORCH_LINUX_WHEEL_VERSION=manylinux_2_28_x86_64

tpu/requirements.txt

@@ -1,13 +1,13 @@
 # TPU Utils
 tpu-info
 # Tensorflow packages
-tensorflow-tpu>=${TENSORFLOW_VERSION}
+tensorflow-tpu==${TENSORFLOW_VERSION}
 --find-links https://storage.googleapis.com/libtpu-tf-releases/index.html
 tensorflow_hub
 tensorflow-io
 tensorflow-probability
 # Torch packages
-torch~=${TORCH_VERSION}
+torch==${TORCH_VERSION}
 https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-${TORCH_VERSION}+libtpu-${PYTHON_WHEEL_VERSION}-${PYTHON_WHEEL_VERSION}-${TORCH_LINUX_WHEEL_VERSION}.whl
 torchaudio==${TORCHAUDIO_VERSION}
 torchvision==${TORCHVISION_VERSION}