Fix jax cuda requirements (#20295)

* fix jax cuda requirements

* fix jax cuda requirements

* fix jax cuda requirements

* fix jax cuda requirements

* fix jax cuda requirements

keras/src/backend/tests/device_scope_test.py

@@ -40,19 +40,19 @@ def test_jax_device_scope(self):
 
         with backend.device_scope("cpu:0"):
             t = backend.numpy.ones((2, 1))
-            self.assertEqual(t.device(), jax.devices("cpu")[0])
+            self.assertEqual(t.device, jax.devices("cpu")[0])
         with backend.device_scope("CPU:0"):
             t = backend.numpy.ones((2, 1))
-            self.assertEqual(t.device(), jax.devices("cpu")[0])
+            self.assertEqual(t.device, jax.devices("cpu")[0])
 
         # When leaving the scope, the device should be back with gpu:0
         t = backend.numpy.ones((2, 1))
-        self.assertEqual(t.device(), jax.devices("gpu")[0])
+        self.assertEqual(t.device, jax.devices("gpu")[0])
 
         # Also verify the explicit gpu device
         with backend.device_scope("gpu:0"):
             t = backend.numpy.ones((2, 1))
-            self.assertEqual(t.device(), jax.devices("gpu")[0])
+            self.assertEqual(t.device, jax.devices("gpu")[0])
 
     @pytest.mark.skipif(backend.backend() != "jax", reason="jax only")
     def test_invalid_jax_device(self):

requirements-jax-cuda.txt

@@ -7,7 +7,9 @@ torch>=2.1.0
 torchvision>=0.16.0
 
 # Jax with cuda support.
-jax[cuda12_pip]
+# TODO: Higher version breaks CI.
+--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+jax[cuda12]==0.4.28
 flax
 
 -r requirements-common.txt