From d7b880d47f4a213b11c99ea28f569bc7ca89c79f Mon Sep 17 00:00:00 2001
From: Divyashree Sreepathihalli <divyashreepathihalli@gmail.com>
Date: Fri, 27 Sep 2024 15:31:38 -0700
Subject: [PATCH] update jax version (#1888)

* update jax version

* remove keras 2 from gpu tests

* update comment
---
 .kokoro/github/ubuntu/gpu/build.sh            | 11 ++----
 .../github/ubuntu/gpu/keras2/continuous.cfg   | 34 -------------------
 .../github/ubuntu/gpu/keras2/presubmit.cfg    | 34 -------------------
 requirements-jax-cuda.txt                     |  4 +--
 4 files changed, 4 insertions(+), 79 deletions(-)
 delete mode 100644 .kokoro/github/ubuntu/gpu/keras2/continuous.cfg
 delete mode 100644 .kokoro/github/ubuntu/gpu/keras2/presubmit.cfg

diff --git a/.kokoro/github/ubuntu/gpu/build.sh b/.kokoro/github/ubuntu/gpu/build.sh
index 6ab351b3f9..ca76446442 100644
--- a/.kokoro/github/ubuntu/gpu/build.sh
+++ b/.kokoro/github/ubuntu/gpu/build.sh
@@ -32,14 +32,7 @@ nvcc --version
 cd "src/github/keras-hub"
 pip install -U pip setuptools psutil
 
-if [ "${KERAS2:-0}" == "1" ]
-then
-   echo "Keras2 detected."
-   pip install -r requirements-common.txt --progress-bar off --timeout 1000
-   pip install tensorflow-text==2.15 tensorflow[and-cuda]~=2.15 keras-core \
-      --timeout 1000
-
-elif [ "$KERAS_BACKEND" == "tensorflow" ]
+if [ "$KERAS_BACKEND" == "tensorflow" ]
 then
    echo "TensorFlow backend detected."
    pip install -r requirements-tensorflow-cuda.txt --progress-bar off \
@@ -67,4 +60,4 @@ then
 else
    pytest keras_hub --check_gpu --run_large \
       --cov=keras-hub
-fi
\ No newline at end of file
+fi
diff --git a/.kokoro/github/ubuntu/gpu/keras2/continuous.cfg b/.kokoro/github/ubuntu/gpu/keras2/continuous.cfg
deleted file mode 100644
index cfceee0de8..0000000000
--- a/.kokoro/github/ubuntu/gpu/keras2/continuous.cfg
+++ /dev/null
@@ -1,34 +0,0 @@
-build_file: "keras-hub/.kokoro/github/ubuntu/gpu/build.sh"
-
-action {
-  define_artifacts {
-    regex: "**/sponge_log.log"
-    regex: "**/sponge_log.xml"
-  }
-}
-
-env_vars: {
-   key: "KERAS2"
-   value: "1"
-}
-
-before_action {
-  fetch_keystore {
-    keystore_resource {
-      keystore_config_id: 73361
-      keyname: "keras_kaggle_username"
-    }
-  }
-}
-
-before_action {
-  fetch_keystore {
-    keystore_resource {
-      keystore_config_id: 73361
-      keyname: "keras_kaggle_secret_key"
-    }
-  }
-}
-
-# Set timeout to 60 mins from default 180 mins
-timeout_mins: 60
\ No newline at end of file
diff --git a/.kokoro/github/ubuntu/gpu/keras2/presubmit.cfg b/.kokoro/github/ubuntu/gpu/keras2/presubmit.cfg
deleted file mode 100644
index cfceee0de8..0000000000
--- a/.kokoro/github/ubuntu/gpu/keras2/presubmit.cfg
+++ /dev/null
@@ -1,34 +0,0 @@
-build_file: "keras-hub/.kokoro/github/ubuntu/gpu/build.sh"
-
-action {
-  define_artifacts {
-    regex: "**/sponge_log.log"
-    regex: "**/sponge_log.xml"
-  }
-}
-
-env_vars: {
-   key: "KERAS2"
-   value: "1"
-}
-
-before_action {
-  fetch_keystore {
-    keystore_resource {
-      keystore_config_id: 73361
-      keyname: "keras_kaggle_username"
-    }
-  }
-}
-
-before_action {
-  fetch_keystore {
-    keystore_resource {
-      keystore_config_id: 73361
-      keyname: "keras_kaggle_secret_key"
-    }
-  }
-}
-
-# Set timeout to 60 mins from default 180 mins
-timeout_mins: 60
\ No newline at end of file
diff --git a/requirements-jax-cuda.txt b/requirements-jax-cuda.txt
index 599a797cc7..222e4400b1 100644
--- a/requirements-jax-cuda.txt
+++ b/requirements-jax-cuda.txt
@@ -8,8 +8,8 @@ torch>=2.1.0
 torchvision>=0.16.0
 
 # Jax with cuda support.
-# TODO: 0.4.24 has an updated Cuda version breaks Jax CI.
+# Keep same version as Keras repo.
 --find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
-jax[cuda12_pip]==0.4.23
+jax[cuda12_pip]==0.4.28
 
 -r requirements-common.txt