Merge pull request #1345 from Kaggle/keras3

Install Keras 3

Author: djherbis

tpu/Dockerfile

@@ -62,6 +62,11 @@ RUN pip install tensorflow_hub https://storage.googleapis.com/cloud-tpu-tpuvm-ar
     kagglehub && \
     /tmp/clean-layer.sh
 
+# b/318672157 Install Keras 3 for loading models.
+# We ignore dependencies for now since tf2.15 should work despite pip saying it won't.
+RUN pip install --no-dependencies "keras>3" keras-cv keras-nlp && \
+    /tmp/clean-layer.sh
+
 # Tensorflow libtpu:
 RUN curl --output /lib/libtpu.so https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/${TF_LIBTPU_VERSION}/libtpu.so
 

tpu/config.txt

@@ -1,9 +1,9 @@
 BASE_IMAGE=python:3.10
 PYTHON_WHEEL_VERSION=cp310
 PYTHON_VERSION_PATH=python3.10
-# https://cloud.google.com/tpu/docs/supported-tpu-configurations#tpu_software_versions:~:text=TensorFlow%20version-,libtpu.so%20version,-2.13.0
-TENSORFLOW_VERSION=2.14.0
-TF_LIBTPU_VERSION=1.8.0
+# https://cloud.google.com/tpu/docs/supported-tpu-configurations#libtpu_versions
+TENSORFLOW_VERSION=2.15.0
+TF_LIBTPU_VERSION=1.9.0
 TF_LINUX_WHEEL_VERSION=manylinux_2_17_x86_64.manylinux2014_x86_64
 JAX_VERSION=0.4.17
 # Supports nightly