code: def __init__(
      self,
      inner_aval: jax_core.AbstractValue,
      memory_space: TPUMemorySpace,
  ):
      assert isinstance(inner_aval, jax_core.ShapedArray)
      self.inner_aval = inner_aval
      self.memory_space = 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 to create a new AbstractMemoryRef object from inherited base classes: state with example: obj = AbstractMemoryRef(inner_aval, memory_space)