Hide JAX's internal tracing state and update libraries to use limited…

… trace-state-querying APIs as needed. This is prep work for stackless which will change those internals while preserving the API.

PiperOrigin-RevId: 677843398

haiku/_src/base.py

@@ -70,12 +70,8 @@ class JaxTraceLevel(NamedTuple):
 
   @classmethod
   def current(cls):
-    # TODO(tomhennigan): Remove once a version of JAX is released incl PR#9423.
-    trace_stack = jax_core.thread_local_state.trace_state.trace_stack.stack
-    top_type = trace_stack[0].trace_type
-    level = trace_stack[-1].level
-    sublevel = jax_core.cur_sublevel()
-    return JaxTraceLevel(opaque=(top_type, level, sublevel))
+    ts = jax_core.get_opaque_trace_state(convention="haiku")
+    return JaxTraceLevel(opaque=ts)
 
 frame_ids = it.count()