update JAX GPU version (#21293)

* test JAX version

* update jax version

* update jax version

* update tensorflow version

* keep tensorflow 2.18.1

* add comment

Author: sachinprasadhs

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

@@ -13,7 +13,8 @@ source venv/bin/activate
 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

@@ -8,7 +8,7 @@ torch==2.6.0
 
 # Jax with cuda support.
 --find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
-jax[cuda12]==0.4.28
+jax[cuda12]==0.6.0
 flax
 
 -r requirements-common.txt