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
@@ -43,13 +44,18 @@ fi
 
 if [ "$KERAS_BACKEND" == "jax" ]
 then
+   export XLA_PYTHON_CLIENT_MEM_FRACTION=.5
+   export JAX_TRACEBACK_FILTERING=off
+
    echo "JAX backend detected."
    pip install -r requirements-jax-cuda.txt --progress-bar off --timeout 1000
    pip uninstall -y keras keras-nightly
    python3 -c 'import jax;print(jax.__version__);print(jax.default_backend())'
    # 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,7 +7,7 @@ action {
   }
 }
 
-env_vars: {
+env_vars {
    key: "KERAS_BACKEND"
    value: "jax"
 }

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

@@ -7,7 +7,7 @@ action {
   }
 }
 
-env_vars: {
+env_vars {
    key: "KERAS_BACKEND"
    value: "jax"
 }