Update Torch and Tensorflow versions in cuda requirements files. (#21732)

Replacement for #21704

Also:
- Disabled an ONNX export test for Torch that was already disabled on GPU with both JAX and Tensorflow.
- Moved install for `tf_keras` from `requirements.txt` to `action.yml` using the `--no-deps` option because `tf_keras` depends on `tensorflow`, which installs the non-CPU version of TensorFlow and causes issues with CPU tests.

Author: hertschuh

.github/workflows/actions.yml

@@ -59,6 +59,7 @@ 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

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,20 @@
 # Tensorflow.
+# Note: when the version of Tensorflow is changed, the version tf_keras must be
+# changed in .github/workflows/actions.yml (pip install --no-deps tf_keras).
 tensorflow-cpu~=2.18.1;sys_platform != 'darwin'
 tensorflow~=2.18.1;sys_platform == 'darwin'
-tf_keras
 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
 flax
+
 # Common deps.
 -r requirements-common.txt