code: def current_jax_trace() -> MainTrace:
      """Returns the innermost Jax tracer."""
      return get_top_trace(())
  
imports: import contextlib
  import dataclasses
  import threading
  import typing
  import jax
  import jax.core
  from jax.core import MainTrace
  from flax.experimental.nnx.nnx import reprlib
label_desc: Use image node_id 2 for calling a global function with example usage: current_jax_trace() and returns: get_top_trace