update jax version (#1888)

* update jax version

* remove keras 2 from gpu tests

* update comment

.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
+fi

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