Update Torch version everywhere and TF version on GPU.

Replacement for #21704

Author: hertschuh

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,21 @@
 # Tensorflow.
 tensorflow-cpu~=2.18.1;sys_platform != 'darwin'
 tensorflow~=2.18.1;sys_platform == 'darwin'
+
+# no-deps otherwise it installs tensorflow (not the CPU version)
+--no-deps
 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-xla==2.6.0;sys_platform != 'darwin'
+torch==2.7.1
+torch-xla==2.7.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