Update libraries to use JAX's limited (and ill-advised) trace-state-querying APIs rather than depending on JAX's deeper internals, which are about to change.

PiperOrigin-RevId: 677843398

Author: dougalm

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")