Fix jax/flax management

Author: fchollet

keras/utils/jax_layer.py

@@ -8,11 +8,7 @@
 from keras.saving import serialization_lib
 from keras.utils import shape_utils
 from keras.utils import tracking
-
-try:
-    import jax
-except ImportError:
-    jax = None
+from keras.utils.module_utils import jax
 
 
 @keras_export("keras.layers.JaxLayer")

keras/utils/module_utils.py

@@ -40,3 +40,4 @@ def __getattr__(self, name):
 gfile = LazyModule("tensorflow.io.gfile", pip_name="tensorflow")
 tensorflow_io = LazyModule("tensorflow_io")
 scipy = LazyModule("scipy")
+jax = LazyModule("jax")

requirements-tensorflow-cuda.txt

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

requirements-torch-cuda.txt

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