From f2ed0424a4359c536a16bb539bee31554ecb4966 Mon Sep 17 00:00:00 2001
From: Richard Stotz <richardstotz@google.com>
Date: Mon, 4 Mar 2024 06:57:30 -0800
Subject: [PATCH] 1.9.0rc0 fixes

PiperOrigin-RevId: 612435124
---
 configure/setup.py                      | 2 +-
 tensorflow_decision_forests/__init__.py | 2 +-
 tools/test_bazel.sh                     | 8 ++++++--
 3 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/configure/setup.py b/configure/setup.py
index f43bb8d..4889f83 100644
--- a/configure/setup.py
+++ b/configure/setup.py
@@ -34,7 +34,7 @@
     "absl_py",
     "wheel",
     "wurlitzer",
-    "tf_keras",
+    "tf_keras~=2.16.0rc2",
 ]
 
 
diff --git a/tensorflow_decision_forests/__init__.py b/tensorflow_decision_forests/__init__.py
index 5ec9ac4..2f11695 100644
--- a/tensorflow_decision_forests/__init__.py
+++ b/tensorflow_decision_forests/__init__.py
@@ -54,7 +54,7 @@
 __version__ = "1.9.0rc0"
 __author__ = "Mathieu Guillame-Bert"
 
-compatible_tf_versions = ["2.16.0rc0"]
+compatible_tf_versions = ["2.16.0-rc0"]
 __git_version__ = "HEAD"  # Modify for release build.
 
 from tensorflow_decision_forests.tensorflow import check_version
diff --git a/tools/test_bazel.sh b/tools/test_bazel.sh
index c13a1bf..452f895 100755
--- a/tools/test_bazel.sh
+++ b/tools/test_bazel.sh
@@ -59,8 +59,12 @@ if [ ${TF_VERSION} == "nightly" ]; then
   TF_MINOR="nightly"
 else
   ${PYTHON} -m pip install tensorflow==${TF_VERSION} --force-reinstall
-  TF_MINOR=$(echo $TF_VERSION | grep -oP '[0-9]+\.[0-9]+')
-  ${PYTHON} -m pip install tf-keras==${TF_MINOR}  --force-reinstall
+  if [[ $TF_VERSION == *"rc"* ]]; then
+    ${PYTHON} -m pip install tf-keras --pre --upgrade
+  else
+    TF_MINOR=$(echo $TF_VERSION | grep -oP '[0-9]+\.[0-9]+')
+    ${PYTHON} -m pip install tf-keras==${TF_MINOR}
+  fi
 fi
 ext=""