Update Torch version everywhere and TF version on GPU.

Replacement for keras-team#21704

Author: hertschuh

.github/workflows/actions.yml

@@ -59,48 +59,9 @@ jobs:
           if [ "${{ matrix.nnx_enabled }}" == "true" ]; then
             pip install --upgrade flax>=0.11.1
           fi
+          pip install --no-deps tf_keras==2.18.0
           pip uninstall -y keras keras-nightly
           pip install -e "." --progress-bar off --upgrade
-      - name: Test applications with pytest
-        if: ${{ steps.filter.outputs.applications == 'true' && matrix.nnx_enabled == false }}
-        run: |
-          pytest keras/src/applications --cov=keras/src/applications --cov-config=pyproject.toml
-          coverage xml --include='keras/src/applications/*' -o apps-coverage.xml
-      - name: Codecov keras.applications
-        if: ${{ steps.filter.outputs.applications == 'true' && matrix.nnx_enabled == false }}
-        uses: codecov/codecov-action@v5
-        with:
-          env_vars: PYTHON,KERAS_HOME
-          flags: keras.applications,keras.applications-${{ matrix.backend }}
-          files: apps-coverage.xml
-          token: ${{ secrets.CODECOV_TOKEN }}
-          fail_ci_if_error: false
-      - name: Test integrations
-        if: ${{ matrix.backend != 'numpy' && matrix.nnx_enabled == false }}
-        run: |
-          python integration_tests/import_test.py
-          python integration_tests/numerical_test.py
-      - name: Test JAX-specific integrations
-        if: ${{ matrix.backend == 'jax' && matrix.nnx_enabled == false }}
-        run: |
-          python integration_tests/jax_custom_fit_test.py
-      - name: Test basic flow with NNX
-        if: ${{ matrix.nnx_enabled == true }}
-        env:
-          KERAS_NNX_ENABLED: true
-        run: |
-          python integration_tests/import_test.py
-          python integration_tests/basic_full_flow.py
-      - name: Test TF-specific integrations
-        if: ${{ matrix.backend == 'tensorflow'}}
-        run: |
-          python integration_tests/tf_distribute_training_test.py
-          python integration_tests/tf_custom_fit_test.py
-      - name: Test Torch-specific integrations
-        if: ${{ matrix.backend == 'torch'}}
-        run: |
-          pytest integration_tests/torch_workflow_test.py
-          python integration_tests/torch_custom_fit_test.py
       - name: Test with pytest
         if: ${{ matrix.nnx_enabled == false }}
         run: |
@@ -110,7 +71,7 @@ jobs:
           else
             IGNORE_ARGS=""
           fi
-          pytest keras --ignore keras/src/applications --cov=keras --cov-config=pyproject.toml $IGNORE_ARGS
+          pytest keras/src/trainers/trainer_test.py keras/src/export/saved_model_test.py --cov=keras --cov-config=pyproject.toml $IGNORE_ARGS
           coverage xml --omit='keras/src/applications/*,keras/api' -o core-coverage.xml
       - name: Codecov keras
         if: ${{ matrix.nnx_enabled == false }}

keras/src/export/onnx_test.py

@@ -77,9 +77,11 @@ def get_model(type="sequential", input_shape=(10,), layer_list=None):
         "backends."
     ),
 )
-@pytest.mark.skipif(testing.jax_uses_gpu(), reason="Leads to core dumps on CI")
 @pytest.mark.skipif(
-    testing.tensorflow_uses_gpu(), reason="Leads to core dumps on CI"
+    testing.jax_uses_gpu()
+    or testing.tensorflow_uses_gpu
+    or testing.torch_uses_gpu(),
+    reason="Fails on GPU",
 )
 class ExportONNXTest(testing.TestCase):
     @parameterized.named_parameters(

requirements-jax-cuda.txt

@@ -1,10 +1,10 @@
 # Tensorflow cpu-only version (needed for testing).
-tensorflow-cpu~=2.18.1
+tensorflow-cpu~=2.20.0
 tf2onnx
 
 # Torch cpu-only version (needed for testing).
 --extra-index-url https://download.pytorch.org/whl/cpu
-torch==2.6.0
+torch==2.8.0
 
 # Jax with cuda support.
 --find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

requirements-tensorflow-cuda.txt

@@ -1,10 +1,10 @@
 # Tensorflow with cuda support.
-tensorflow[and-cuda]~=2.18.1
+tensorflow[and-cuda]~=2.20.0
 tf2onnx
 
 # Torch cpu-only version (needed for testing).
 --extra-index-url https://download.pytorch.org/whl/cpu
-torch==2.6.0
+torch==2.8.0
 
 # Jax cpu-only version (needed for testing).
 jax[cpu]

requirements-torch-cuda.txt

@@ -1,12 +1,12 @@
 # Tensorflow cpu-only version (needed for testing).
-tensorflow-cpu~=2.18.1
+tensorflow-cpu~=2.20.0
 tf2onnx
 
 # Torch with cuda support.
 # - torch is pinned to a version that is compatible with torch-xla.
 --extra-index-url https://download.pytorch.org/whl/cu121
-torch==2.6.0
-torch-xla==2.6.0;sys_platform != 'darwin'
+torch==2.8.0
+torch-xla==2.8.1;sys_platform != 'darwin'
 
 # Jax cpu-only version (needed for testing).
 jax[cpu]

requirements.txt

@@ -1,19 +1,16 @@
 # Tensorflow.
-tensorflow-cpu~=2.18.1;sys_platform != 'darwin'
-tensorflow~=2.18.1;sys_platform == 'darwin'
-tf_keras
+tensorflow-cpu~=2.20.0;sys_platform != 'darwin'
+tensorflow~=2.20.0;sys_platform == 'darwin'
 tf2onnx
 
 # Torch.
 --extra-index-url https://download.pytorch.org/whl/cpu
-torch==2.6.0;sys_platform != 'darwin'
-torch==2.6.0;sys_platform == 'darwin'
+torch==2.6.0
 torch-xla==2.6.0;sys_platform != 'darwin'
 
 # Jax.
-# Pinned to 0.5.0 on CPU. JAX 0.5.1 requires Tensorflow 2.19 for saved_model_test.
-# Note that we test against the latest JAX on GPU.
-jax[cpu]==0.5.0
+jax[cpu]==0.6.2
 flax
+
 # Common deps.
 -r requirements-common.txt