code: def _ref_raise_to_shaped(ref_aval: AbstractMemoryRef, weak_type):
      return AbstractMemoryRef(
          jax_core.raise_to_shaped(ref_aval.inner_aval, weak_type),
          ref_aval.memory_space,
      )
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import dataclasses
  import enum
  import functools
  from typing import Any
  from jax._src import core
  from jax._src import dtypes
  from jax._src import state
  from jax._src import tree_util
  from jax._src import util
  import jax.numpy
  from jax._src.pallas import core
label_desc: Use image node_id 1 for calling a global function with example usage: _ref_raise_to_shaped(ref_aval, weak_type) and returns: AbstractMemoryRef