code: def nnx_trace(trace: MainTrace):
      TRACE_CONTEXT.nnx_trace_stack.append(trace)
      try:
          yield
      finally:
          TRACE_CONTEXT.nnx_trace_stack.pop()
  
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 5 for calling a global function with example usage: nnx_trace(trace) without return types