In open source tests, use tf-keras-nightly (Keras 2) instead of Keras 3.

PiperOrigin-RevId: 575910204
tensorflow · Oct 23, 2023 · abdc777 · abdc777
1 parent 73a4a75
commit abdc777
Showing 1 changed file with 3 additions and 1 deletion.
diff --git 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() {