Stackless yashful

PiperOrigin-RevId: 681582933
google · Oct 4, 2024 · c9572a1 · c9572a1
1 parent 36fd2e1
commit c9572a1
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/learned_optimization/jax_utils.py b/learned_optimization/jax_utils.py
@@ -46,9 +46,12 @@ def body_fn(_, operand):
 
 def in_jit() -> bool:
   """Returns true if tracing jit."""
-  return "DynamicJaxprTrace" in str(
-      jax.core.thread_local_state.trace_state.trace_stack
-  )
+  if jax.__version_info__ <= (0, 4, 33):
+    return "DynamicJaxprTrace" in str(
+        jax.core.thread_local_state.trace_state.trace_stack  # type: ignore
+    )
+
+  return jax.core.unsafe_am_i_under_a_jit_DO_NOT_USE()
 
 
 Carry = TypeVar("Carry")