Fix JAX GPU tests.

Added debugging.

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,25 +44,31 @@ fi
 
 if [ "$KERAS_BACKEND" == "jax" ]
 then
+   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.
-   pytest keras --ignore keras/src/applications \
-               --ignore keras/src/layers/merging/merging_test.py \
-               --ignore keras/src/trainers/data_adapters/py_dataset_adapter_test.py \
-               --ignore keras/src/backend/jax/distribution_lib_test.py \
-               --ignore keras/src/distribution/distribution_lib_test.py \
-               --cov=keras \
-               --cov-config=pyproject.toml
+   # pytest keras --ignore keras/src/applications \
+   #             --ignore keras/src/layers/merging/merging_test.py \
+   #             --ignore keras/src/trainers/data_adapters/py_dataset_adapter_test.py \
+   #             --ignore keras/src/backend/jax/distribution_lib_test.py \
+   #             --ignore keras/src/distribution/distribution_lib_test.py \
+   #             --cov=keras \
+   #             --cov-config=pyproject.toml
+
+   pytest -s keras/src/utils/jax_layer_test.py
 
-   pytest keras/src/distribution/distribution_lib_test.py --cov=keras --cov-config=pyproject.toml
+   # pytest keras/src/distribution/distribution_lib_test.py --cov=keras --cov-config=pyproject.toml
 fi
 
 if [ "$KERAS_BACKEND" == "torch" ]

conftest.py

@@ -1,5 +1,7 @@
 import os
 
+from absl import logging
+
 # When using jax.experimental.enable_x64 in unit test, we want to keep the
 # default dtype with 32 bits, aligning it with Keras's default.
 os.environ["JAX_DEFAULT_DTYPE_BITS"] = "32"
@@ -18,6 +20,9 @@
 
 
 def pytest_configure(config):
+    logging.use_python_logging()
+    logging.set_verbosity(5)
+
     config.addinivalue_line(
         "markers",
         "requires_trainable_backend: mark test for trainable backend only",

requirements-jax-cuda.txt

@@ -1,6 +1,6 @@
 # Tensorflow cpu-only version (needed for testing).
-tensorflow-cpu~=2.18.1
-tf2onnx
+tensorflow-cpu~=2.19.0
+# tf2onnx
 
 # Torch cpu-only version (needed for testing).
 --extra-index-url https://download.pytorch.org/whl/cpu