Fix JAX GPU tests.

Using XLA_PYTHON_CLIENT_MEM_FRACTION.

Author: hertschuh

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

@@ -18,6 +18,7 @@ python3 --version
 # Check cuda
 nvidia-smi
 nvcc --version
+echo "LD_LIBRARY_PATH before ${LD_LIBRARY_PATH}"
 
 cd "src/github/keras"
 pip install -U pip setuptools
@@ -50,6 +51,8 @@ then
    # Raise error if GPU is not detected.
    python3 -c 'import jax;assert jax.default_backend().lower() == "gpu"'
 
+   echo "LD_LIBRARY_PATH after ${LD_LIBRARY_PATH}"
+
    # TODO: keras/layers/merging/merging_test.py::MergingLayersTest::test_sparse_dot_2d Fatal Python error: Aborted
    # TODO: keras/trainers/data_adapters/py_dataset_adapter_test.py::PyDatasetAdapterTest::test_basic_flow0 Fatal Python error: Aborted
    # keras/backend/jax/distribution_lib_test.py is configured for CPU test for now.

.kokoro/github/ubuntu/gpu/jax/continuous.cfg

@@ -7,10 +7,20 @@ action {
   }
 }
 
-env_vars: {
+env_vars {
    key: "KERAS_BACKEND"
    value: "jax"
 }
 
+env_vars {
+   key: "XLA_PYTHON_CLIENT_MEM_FRACTION"
+   value: "0.5"
+}
+
+env_vars {
+   key: "JAX_TRACEBACK_FILTERING"
+   value: "off"
+}
+
 # Set timeout to 60 mins from default 180 mins
 timeout_mins: 60

.kokoro/github/ubuntu/gpu/jax/presubmit.cfg

@@ -7,10 +7,20 @@ action {
   }
 }
 
-env_vars: {
+env_vars {
    key: "KERAS_BACKEND"
    value: "jax"
 }
 
+env_vars {
+   key: "XLA_PYTHON_CLIENT_MEM_FRACTION"
+   value: "0.5"
+}
+
+env_vars {
+   key: "JAX_TRACEBACK_FILTERING"
+   value: "off"
+}
+
 # Set timeout to 60 mins from default 180 mins
 timeout_mins: 60