From abdc77761a074a35def9b6b08927952c19b51a76 Mon Sep 17 00:00:00 2001
From: jburnim <jburnim@google.com>
Date: Mon, 23 Oct 2023 13:20:37 -0700
Subject: [PATCH] In open source tests, use tf-keras-nightly (Keras 2) instead
 of Keras 3.

PiperOrigin-RevId: 575910204
---
 testing/dependency_install_lib.sh | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/testing/dependency_install_lib.sh b/testing/dependency_install_lib.sh
index 9db1691815..801d7a3361 100644
--- a/testing/dependency_install_lib.sh
+++ b/testing/dependency_install_lib.sh
@@ -69,7 +69,9 @@ install_tensorflow() {
   PIP_FLAGS=${2-}
   # NB: tf-nightly pulls in other deps, like numpy, absl, and six, transitively.
   TF_VERSION_STR=$(find_good_tf_nightly_version_str $TF_NIGHTLY_PACKAGE)
-  python -m pip install $PIP_FLAGS $TF_NIGHTLY_PACKAGE==$TF_VERSION_STR
+  python -m pip install $PIP_FLAGS \
+    $TF_NIGHTLY_PACKAGE==$TF_VERSION_STR \
+    tf-keras-nightly
 }
 
 install_jax() {