Update libraries to use JAX's limited (and ill-advised) trace-state-q…

…uerying APIs rather than depending on JAX's deeper internals, which are about to change. PiperOrigin-RevId: 678332660
google · Sep 24, 2024 · 143b3db · 143b3db
1 parent 4bcaeb0
commit 143b3db
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
+    )
+
+  return jax.core.unsafe_am_i_under_a_jit_DO_NOT_USE()
 
 
 Carry = TypeVar("Carry")