Skip to content

Commit 3d5ba66

Browse files
authored
Revert "Keras 3 w/ tf2.16" (#1387)
Reverts #1385 Reverts #1384 tf/libtpu combo is crashing, will have to wait to upgrade
1 parent 1369652 commit 3d5ba66

File tree

2 files changed

+7
-3
lines changed

2 files changed

+7
-3
lines changed

tpu/Dockerfile

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,10 +65,14 @@ RUN pip install tensorflow_hub https://storage.googleapis.com/cloud-tpu-tpuvm-ar
6565
papermill jupyterlab python-lsp-server[all] "jupyter-lsp==1.5.1" \
6666
pandas matplotlib opencv-python-headless librosa accelerate diffusers scikit-learn transformers \
6767
seaborn timm albumentations einops pyarrow fastparquet opencv-python \
68-
"keras>3" keras-cv keras-nlp \
6968
kagglehub && \
7069
/tmp/clean-layer.sh
7170

71+
# b/318672157 Install Keras 3 for loading models.
72+
# We ignore dependencies for now since tf2.15 should work despite pip saying it won't.
73+
RUN pip install --no-dependencies "keras>3" keras-cv keras-nlp && \
74+
/tmp/clean-layer.sh
75+
7276
# Tensorflow libtpu:
7377
RUN curl --output /lib/libtpu.so https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/${TF_LIBTPU_VERSION}/libtpu.so
7478

tpu/config.txt

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,8 @@ PYTHON_WHEEL_VERSION=cp310
33
PYTHON_VERSION_PATH=python3.10
44
# gsutil ls gs://cloud-tpu-tpuvm-artifacts/tensorflow
55
# https://cloud.google.com/tpu/docs/supported-tpu-configurations#libtpu_versions
6-
TENSORFLOW_VERSION=2.16.1
7-
TF_LIBTPU_VERSION=1.10.1
6+
TENSORFLOW_VERSION=2.15.0
7+
TF_LIBTPU_VERSION=1.9.0
88
TF_LINUX_WHEEL_VERSION=manylinux_2_17_x86_64.manylinux2014_x86_64
99
JAX_VERSION=0.4.26
1010
# gsutil ls gs://pytorch-xla-releases/wheels/tpuvm/* | grep libtpu | grep -v -E ".*rc[0-9].*"

0 commit comments

Comments
 (0)