code: def run_state_reference(f: Callable[..., None]):
      def wrapped(args):
          flat_args, in_tree = tree_util.tree_flatten(args)
          avals = [
              core.raise_to_shaped(core.get_aval(arg))
              for arg in flat_args
          ]
          jaxpr_, consts, _ = initial_style_jaxpr(
              f, in_tree, map(AbstractRef, avals)
          )
          jaxpr = hoist_consts_to_refs(jaxpr_)
          discharged_jaxpr, discharged_consts = discharge_state(
              jaxpr, ()
          )
          out_const_flat = core.eval_jaxpr(
              discharged_jaxpr, discharged_consts, *consts, *args
          )
          _, out_flat = split_list(out_const_flat, [len(consts)])
          return in_tree.unflatten(out_flat)
  
      return wrapped
  
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 33 for calling a global function with example usage: run_state_reference(f) and returns: wrapped, in_tree