Fix gpu ci by adding flax in the requirements (#19349)

Author: james77777778

keras/utils/jax_layer_test.py

@@ -420,6 +420,7 @@ def test_jax_layer(
             "non_trainable_params": 536,
         },
     )
+    @pytest.mark.skipif(flax is None, reason="Flax library is not available.")
     def test_flax_layer(
         self,
         flax_model_class,
@@ -575,6 +576,7 @@ def jax_fn(params, inputs):
         test_output = model(test_inputs)
         self.assertAllClose(test_output, np.ones((2, 60, 3)))
 
+    @pytest.mark.skipif(flax is None, reason="Flax library is not available.")
     def test_with_flax_state_no_params(self):
         class MyFlaxLayer(flax.linen.Module):
             @flax.linen.compact

requirements-jax-cuda.txt

@@ -10,5 +10,6 @@ torchvision>=0.16.0
 # TODO: 0.4.24 has an updated Cuda version breaks Jax CI.
 --find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 jax[cuda12_pip]==0.4.23
+flax
 
 -r requirements-common.txt

requirements-tensorflow-cuda.txt

@@ -8,5 +8,6 @@ torchvision>=0.16.0
 
 # Jax cpu-only version (needed for testing).
 jax[cpu]
+flax
 
 -r requirements-common.txt

requirements-torch-cuda.txt

@@ -8,5 +8,6 @@ torchvision==0.17.1+cu121
 
 # Jax cpu-only version (needed for testing).
 jax[cpu]
+flax
 
 -r requirements-common.txt