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

flax/core/tracers.py

@@ -20,15 +20,12 @@
 
 
 def current_trace():
-  """Returns the innermost Jax tracer."""
-  return jax.core.find_top_trace(())
+  """Returns the current JAX state tracer."""
+  return jax.core.get_opaque_trace_state(convention="flax")
 
 
 def trace_level(main):
-  """Returns the level of the trace of -infinity if it is None."""
-  if main:
-    return main.level
-  return float('-inf')
+  return main
 
 
 def check_trace_level(base_level):

flax/nnx/tracers.py

@@ -15,14 +15,14 @@
 # Taken from flax/core/tracer.py 🏴‍☠️
 
 
-from jax.core import MainTrace, thread_local_state
+from jax.core import get_opaque_trace_state, OpaqueTraceState
 
 from flax.nnx import reprlib
 
 
-def current_jax_trace() -> MainTrace:
-  """Returns the innermost Jax tracer."""
-  return thread_local_state.trace_state.trace_stack.dynamic
+def current_jax_trace() -> OpaqueTraceState:
+  """Returns the Jax tracing state."""
+  return get_opaque_trace_state(convention="nnx")
 
 
 class TraceState(reprlib.Representable):
@@ -36,7 +36,7 @@ def jax_trace(self):
     return self._jax_trace
 
   def is_valid(self) -> bool:
-    return self._jax_trace is current_jax_trace()
+    return self._jax_trace == current_jax_trace()
 
   def __nnx_repr__(self):
     yield reprlib.Object(f'{type(self).__name__}')
@@ -52,4 +52,4 @@ def __treescope_repr__(self, path, subtree_renderer):
     )
 
   def __eq__(self, other):
-    return isinstance(other, TraceState) and self._jax_trace is other._jax_trace
+    return isinstance(other, TraceState) and self._jax_trace == other._jax_trace