code: def get_top_trace(pytree: tp.Union[tp.Any, Tracer]) -> MainTrace:
      """Returns the main top trace of a sequence of tracers."""
      if isinstance(pytree, Tracer):
          return pytree._trace.main
  
      return jax.core.find_top_trace(
          jax.tree_util.tree_leaves(pytree)
      ).main
  
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 1 for calling a global function with example usage: get_top_trace(pytree) and returns: jax, pytree