Update requirements-jax-cuda.txt (keras-team#2252)

* Update requirements-jax-cuda.txt

Upgrading jax version. For CI tests for keras-hub on jax GPU

* Update build.sh

For CI tests for JAX GPU

* Update build.sh

* Update build.sh

* Update requirements-jax-cuda.txt

* Update requirements-jax-cuda.txt

* Update requirements-jax-cuda.txt

* Update requirements-jax-cuda.txt

* Update build.sh

* Update build.sh

* Update build.sh

Author: pctablet505

.kokoro/github/ubuntu/gpu/build.sh

@@ -16,15 +16,16 @@ fi
 set -x
 cd "${KOKORO_ROOT}/"
 
-PYTHON_BINARY="/usr/bin/python3.9"
+PYTHON_BINARY="/usr/bin/python3.10"
 
 "${PYTHON_BINARY}" -m venv venv
 source venv/bin/activate
 # Check the python version
 python --version
 python3 --version
 
-export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64:"
+# setting the LD_LIBRARY_PATH manually is causing segmentation fault
+# export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64:"
 # Check cuda
 nvidia-smi
 nvcc --version

requirements-jax-cuda.txt

@@ -10,6 +10,6 @@ torchvision>=0.16.0
 # Jax with cuda support.
 # Keep same version as Keras repo.
 --find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
-jax[cuda12_pip]==0.4.28
+jax[cuda12]==0.6.0
 
 -r requirements-common.txt