Upgrade TPU image to Python 3.11.Tpupy311 (#1493)

Python 3.10 is entering its [last year of
support](https://devguide.python.org/versions/) before end-of-life and
many packages, including [NumPy](https://devguide.python.org/versions/),
have dropped support for it altogether.

Included in this change:

* Upgrade the TPU docker image to derive from `python:3.11`.
* Upgrade `tensorflow` to 2.20.0.
* Upgrade `jax` to >=0.5.2. For a compatible dep closure, this installs
`jax` 0.7.2 re: `tensorflow-tpu` dep on`libtpu`.
* Upgrade `torch` (and ecosystem) to 2.8.0. Of note, there is no wheel
with a `+libtpu` label.
* Remove unneeded environment variable.

Tested:

Locally, by invoking `./tpu/build`:

> <img width="1116" height="297" alt="9jyLVT6hAPjKaCh"
src="https://github.com/user-attachments/assets/e1c8e37e-5e65-4f43-807f-4ffca8d6b6ac"
/>

Also invoked other back-end testing.

Author: metrizable

tpu/Dockerfile

@@ -55,7 +55,6 @@ RUN sed -i '/from tensorflow_hub import uncompressed_module_resolver/a from tens
 RUN sed -i '/_install_default_resolvers()/a \ \ registry.resolver.add_implementation(kaggle_module_resolver.KaggleFileResolver())' /usr/local/lib/${PYTHON_VERSION_PATH}/site-packages/tensorflow_hub/config.py
 
 # Set these env vars so that they don't produce errs calling the metadata server to load them:
-ENV TPU_ACCELERATOR_TYPE=v3-8
 ENV TPU_PROCESS_ADDRESSES=local
 
 # Metadata

tpu/config.txt

@@ -1,12 +1,12 @@
-BASE_IMAGE=python:3.10
-PYTHON_WHEEL_VERSION=cp310
-PYTHON_VERSION_PATH=python3.10
-TENSORFLOW_VERSION=2.18.0
+BASE_IMAGE=python:3.11
+PYTHON_WHEEL_VERSION=cp311
+PYTHON_VERSION_PATH=python3.11
+TENSORFLOW_VERSION=2.20.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.6.0
+TORCH_VERSION=2.8.0
 # https://github.com/pytorch/audio supports nightly
-TORCHAUDIO_VERSION=2.6.0
+TORCHAUDIO_VERSION=2.8.0
 # https://github.com/pytorch/vision supports nightly
-TORCHVISION_VERSION=0.21.0
+TORCHVISION_VERSION=0.23.0
 TORCH_LINUX_WHEEL_VERSION=manylinux_2_28_x86_64

tpu/requirements.txt

@@ -8,11 +8,11 @@ tensorflow-io
 tensorflow-probability
 # Torch packages
 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
+https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-${TORCH_VERSION}-${PYTHON_WHEEL_VERSION}-${PYTHON_WHEEL_VERSION}-${TORCH_LINUX_WHEEL_VERSION}.whl
 torchaudio==${TORCHAUDIO_VERSION}
 torchvision==${TORCHVISION_VERSION}
 # Jax packages
-jax[tpu]>=0.4.34
+jax[tpu]>=0.5.2
 --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
 distrax
 flax