code: def val_to_ref_aval(x) -> AbstractRef:
      aval = core.raise_to_shaped(core.get_aval(x))
      if type(aval) is not core.ShapedArray:
          raise Exception(f"can't make ref from {x}")
      return AbstractRef(aval)
  
file_path: jax/jax/_src/state/utils.py
imports: from collections.abc import Sequence
  from jax.interpreters import partial_eval
  from jax._src import core
  from jax._src import linear_util
  from jax._src.state import AbstractRef
  from jax._src.util import partition_list, merge_lists, split_list, safe_map, safe_zip
  from jax._src.state.primitives import ref_get
label_desc: Use image node_id 2 for calling a global function with example usage: val_to_ref_aval(x) and returns: AbstractRef