diff --git a/tpu/Dockerfile b/tpu/Dockerfile
index fe79260f..b11bb59d 100644
--- a/tpu/Dockerfile
+++ b/tpu/Dockerfile
@@ -65,10 +65,14 @@ RUN pip install tensorflow_hub https://storage.googleapis.com/cloud-tpu-tpuvm-ar
     papermill jupyterlab python-lsp-server[all] "jupyter-lsp==1.5.1" \
     pandas matplotlib opencv-python-headless librosa accelerate diffusers scikit-learn transformers \
     seaborn timm albumentations einops pyarrow fastparquet opencv-python \
-    "keras>3" keras-cv keras-nlp \
     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
 
diff --git a/tpu/config.txt b/tpu/config.txt
index 1f21961c..c893518f 100644
--- a/tpu/config.txt
+++ b/tpu/config.txt
@@ -3,8 +3,8 @@ PYTHON_WHEEL_VERSION=cp310
 PYTHON_VERSION_PATH=python3.10
 # gsutil ls gs://cloud-tpu-tpuvm-artifacts/tensorflow
 # https://cloud.google.com/tpu/docs/supported-tpu-configurations#libtpu_versions
-TENSORFLOW_VERSION=2.16.1
-TF_LIBTPU_VERSION=1.10.1
+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.26
 # gsutil ls gs://pytorch-xla-releases/wheels/tpuvm/* | grep libtpu | grep -v -E ".*rc[0-9].*"