code: def _initial_style_jaxpr(fun, in_tree, in_avals):
      fun_, out_tree_thunk = api_util.flatten_fun_nokwargs(
          lu.wrap_init(fun), tree_util.treedef_tuple((in_tree,))
      )
      debug = pe.debug_info(
          fun_, in_tree, out_tree_thunk, False, "run_state"
      )
      jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(
          fun_, in_avals, debug
      )
      return jaxpr, consts, out_tree_thunk()
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import dataclasses
  from functools import partial
  import operator
  from typing import Any, Callable, Protocol
  import numpy
  from jax._src import api_util
  from jax._src import ad_util
  from jax._src import config
  from jax._src import core
  from jax._src import linear_util
  from jax._src import source_info_util
  from jax._src import tree_util
  from jax._src.interpreters import ad
  from jax._src.interpreters import mlir
  from jax._src.interpreters import partial_eval
  from jax._src.lax import lax
  from jax._src.lax import slicing
  from jax._src.state import indexing
  from jax._src.state.types import AbstractRef, RefEffect
  from jax._src.state.primitives import get_p, swap_p, addupdate_p
  from jax._src.state.utils import hoist_consts_to_refs
  from jax._src.typing import Array
  from jax._src.util import safe_map, safe_zip, split_list, weakref_lru_cache, partition_list, merge_lists, split_dict
label_desc: Use image node_id 31 for calling a global function with example usage: _initial_style_jaxpr(fun, in_tree, in_avals) and returns: jaxpr,  consts,  out_tree_thunk