From f16b4537f18a0640c119f3ee27b37b06284ec53a Mon Sep 17 00:00:00 2001
From: Vincent Roseberry <rosbo@google.com>
Date: Fri, 19 Jan 2024 03:58:24 +0000
Subject: [PATCH 1/5] [WIP] Upgrade to Keras 3

Work in progress, waiting for other tensorflow related packages to be
upgraded. See linked bug for status.

http://b/318672158
---
 Dockerfile.tmpl | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/Dockerfile.tmpl b/Dockerfile.tmpl
index 245228df..3a7ba5af 100644
--- a/Dockerfile.tmpl
+++ b/Dockerfile.tmpl
@@ -206,14 +206,17 @@ RUN pip install \
         "tensorflow==${TENSORFLOW_VERSION}" \
         "tensorflow-io==${TENSORFLOW_IO_VERSION}" \
         tensorflow_decision_forests \
+        tensorflow-probability \
         tensorflow_text \
         tensorflowjs \
         tensorflow_hub && \
     /tmp/clean-layer.sh
 
-# TODO(b/318672158): Upgrade to Keras 3 once compatible with other TF libries.
-# See blockers here: https://b.corp.google.com/issues/319722433#comment8
-RUN pip install keras keras-cv keras-nlp && \
+# Remove --no-deps flag once Keras 3.* is included in our base image.
+# We ignore dependencies since tf2.15 and Keras 3.* should work despite pip saying it won't.
+# Currently, keras tries to install a nightly version of tf 2.16: https://github.com/keras-team/keras/blob/fe2f54aa5bc42fb23a96449cf90434ab9bb6a2cd/requirements.txt#L2
+RUN pip install --no-deps keras>= 3.0.2 && \
+    pip install keras-cv keras-nlp && \
     /tmp/clean-layer.sh
 
 RUN pip install pysal

From 7f630e8824b195d3ade0d9ff9acc082a0c059515 Mon Sep 17 00:00:00 2001
From: Vincent Roseberry <rosbo@google.com>
Date: Fri, 19 Jan 2024 18:41:30 +0000
Subject: [PATCH 2/5] Fix keras install command

---
 Dockerfile.tmpl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Dockerfile.tmpl b/Dockerfile.tmpl
index 3a7ba5af..999baf65 100644
--- a/Dockerfile.tmpl
+++ b/Dockerfile.tmpl
@@ -215,7 +215,7 @@ RUN pip install \
 # Remove --no-deps flag once Keras 3.* is included in our base image.
 # We ignore dependencies since tf2.15 and Keras 3.* should work despite pip saying it won't.
 # Currently, keras tries to install a nightly version of tf 2.16: https://github.com/keras-team/keras/blob/fe2f54aa5bc42fb23a96449cf90434ab9bb6a2cd/requirements.txt#L2
-RUN pip install --no-deps keras>= 3.0.2 && \
+RUN pip install --no-deps -U keras && \
     pip install keras-cv keras-nlp && \
     /tmp/clean-layer.sh
 

From a23e21cfbab31553ee2f7622d7315d123ab4e086 Mon Sep 17 00:00:00 2001
From: Vincent Roseberry <rosbo@google.com>
Date: Fri, 19 Jan 2024 21:48:58 +0000
Subject: [PATCH 3/5] avoid keras downgrade by keras-nlp

---
 Dockerfile.tmpl | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/Dockerfile.tmpl b/Dockerfile.tmpl
index 999baf65..9286eafc 100644
--- a/Dockerfile.tmpl
+++ b/Dockerfile.tmpl
@@ -212,11 +212,10 @@ RUN pip install \
         tensorflow_hub && \
     /tmp/clean-layer.sh
 
-# Remove --no-deps flag once Keras 3.* is included in our base image.
+# Remove "--no-deps" flag and "namex" package once Keras 3.* is included in our base image.
 # We ignore dependencies since tf2.15 and Keras 3.* should work despite pip saying it won't.
 # Currently, keras tries to install a nightly version of tf 2.16: https://github.com/keras-team/keras/blob/fe2f54aa5bc42fb23a96449cf90434ab9bb6a2cd/requirements.txt#L2
-RUN pip install --no-deps -U keras && \
-    pip install keras-cv keras-nlp && \
+RUN pip install --no-deps "keras>3" keras-cv keras-nlp namex && \
     /tmp/clean-layer.sh
 
 RUN pip install pysal

From fc9eb8fde739e9749e9a5347c3e61bca3d2f85fa Mon Sep 17 00:00:00 2001
From: Vincent Roseberry <rosbo@google.com>
Date: Fri, 26 Jan 2024 00:57:07 +0000
Subject: [PATCH 4/5] Install tf-keras and use unreleased tfp version

---
 Dockerfile.tmpl | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/Dockerfile.tmpl b/Dockerfile.tmpl
index 9286eafc..9fd492d5 100644
--- a/Dockerfile.tmpl
+++ b/Dockerfile.tmpl
@@ -202,14 +202,16 @@ RUN apt-get install -y default-jre && \
 
 RUN pip install -f http://h2o-release.s3.amazonaws.com/h2o/latest_stable_Py.html h2o && /tmp/clean-layer.sh
 
+# b/318672158 Use simply tensorflow-probability once > 0.23.0 is released.
 RUN pip install \
         "tensorflow==${TENSORFLOW_VERSION}" \
         "tensorflow-io==${TENSORFLOW_IO_VERSION}" \
         tensorflow_decision_forests \
-        tensorflow-probability \
+        git+https://github.com/tensorflow/probability.git@fbc5ebe9b1d343113fb917010096cfd88b32eecf \
         tensorflow_text \
         tensorflowjs \
-        tensorflow_hub && \
+        tensorflow_hub \
+        tf-keras && \
     /tmp/clean-layer.sh
 
 # Remove "--no-deps" flag and "namex" package once Keras 3.* is included in our base image.

From fc4e4042b3fa22209d372e0074331d2adb61d00b Mon Sep 17 00:00:00 2001
From: Vincent Roseberry <rosbo@google.com>
Date: Fri, 26 Jan 2024 17:46:25 +0000
Subject: [PATCH 5/5] upgrade tensorflow_hub

---
 Dockerfile.tmpl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Dockerfile.tmpl b/Dockerfile.tmpl
index 9fd492d5..e9f9364b 100644
--- a/Dockerfile.tmpl
+++ b/Dockerfile.tmpl
@@ -210,7 +210,7 @@ RUN pip install \
         git+https://github.com/tensorflow/probability.git@fbc5ebe9b1d343113fb917010096cfd88b32eecf \
         tensorflow_text \
         tensorflowjs \
-        tensorflow_hub \
+        "tensorflow_hub>=0.16.0" \
         tf-keras && \
     /tmp/clean-layer.sh