code: def _run_state_abstract_eval(
      *avals: core.AbstractValue,
      jaxpr: core.Jaxpr,
      which_linear: tuple[bool, ...],
  ):
      del which_linear
      # When we abstractly evaluate `run_state`, we want to keep track of which
      # input avals are `Ref`s and which are not. If an aval is a `Ref`, we want to
      # "propagate" out its inner effects. Otherwise, the effects are local to this
      # `run_state`.
      is_ref = {
          i
          for i, aval in enumerate(avals)
          if isinstance(aval, AbstractRef)
      }
      nonlocal_effects = {
          e
          for e in jaxpr.effects
          if (isinstance(e, RefEffect) and e.input_index in is_ref)
          or not isinstance(e, RefEffect)
      }
      return avals, nonlocal_effects
  
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 21 for calling a global function with example usage: _run_state_abstract_eval() and returns: avals,  nonlocal_effects