Move Jax import to conditional

keras-team · Mar 21, 2024 · 9ceb80b · 9ceb80b
1 parent dc398f7
commit 9ceb80b
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/keras/utils/jax_layer.py b/keras/utils/jax_layer.py
@@ -1,6 +1,5 @@
 import inspect
 
-import jax
 import numpy as np
 
 from keras import backend
@@ -10,6 +9,11 @@
 from keras.utils import shape_utils
 from keras.utils import tracking
 
+try:
+    import jax
+except ImportError:
+    jax = None
+
 
 @keras_export("keras.layers.JaxLayer")
 class JaxLayer(Layer):