code: def to_graph(fun):
      """Converts a Haiku function into an graph IR (extracted for testing)."""
  
      @functools.wraps(fun)
      def wrapped_fun(*args):
          """See `fun`."""
          f = lu.wrap_init(fun)
          args_flat, in_tree = jax.tree_util.tree_flatten((args, {}))
          flat_fun, out_tree = jax.api_util.flatten_fun(f, in_tree)
          graph = Graph.create(title=name_or_str(fun))
  
          @contextlib.contextmanager
          def method_hook(mod: module.Module, method_name: str):
              subg = Graph.create()
              with graph_stack(subg):
                  yield
              title = mod.module_name
              if method_name != "__call__":
                  title += f" ({method_name})"
              graph_stack.peek().subgraphs.append(
                  subg.evolve(title=title)
              )
  
          with graph_stack(graph), module.hook_methods(
              method_hook
          ), jax.core.new_main(DotTrace) as main:
              out_flat = _interpret_subtrace(
                  flat_fun, main
              ).call_wrapped(*args_flat)
          out = jax.tree_util.tree_unflatten(out_tree(), out_flat)
  
          return graph, args, out
  
      return wrapped_fun
  
imports: import collections
  import contextlib
  import functools
  import html
  from typing import Any, Callable, NamedTuple, Optional
  from haiku._src import data_structures
  from haiku._src import module
  from haiku._src import utils
  import jax
  import jax.core
  from jax.experimental import pjit
  from jax.extend import linear_util
label_desc: Use image node_id 4 for calling a global function with example usage: to_graph(fun) and returns: wrapped_fun, graph,  args,  out