Add jax_utils

keras-team · Aug 31, 2023 · 91a58b7 · 91a58b7
1 parent 14b3755
commit 91a58b7
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/keras_core/utils/jax_utils.py b/keras_core/utils/jax_utils.py
@@ -0,0 +1,9 @@
+from keras_core import backend
+
+
+def is_in_jax_tracing_scope():
+    if backend.backend() == "jax":
+        x = backend.numpy.ones(())
+        if x.__class__.__name__ == "DynamicJaxprTracer":
+            return True
+    return False