[Pallas/Hijax] Let MemoryRef be parameterized by a type, instead of forcing it to be ShapedArray

This will allow us to generalize to types beyond ShapedArray (new HiTypes specifically)

PiperOrigin-RevId: 807400923

Author: sharadmv

jax/_src/pallas/core.py

@@ -210,24 +210,34 @@ def update(
 @dataclasses.dataclass(frozen=True)
 class MemoryRef:
   """Like jax.ShapeDtypeStruct but with memory spaces."""
-  shape: tuple[int, ...]
-  dtype: jnp.dtype | dtypes.ExtendedDType
+  inner_aval: jax_core.AbstractValue
   # TODO(b/368122763): Unify memory space types across backends
   memory_space: Any
 
   def get_array_aval(self) -> jax_core.ShapedArray:
-    dtype = self.dtype
+    if not isinstance(self.inner_aval, jax_core.ShapedArray):
+      raise ValueError(
+          f"MemoryRef type must be a ShapedArray, got {type(self.inner_aval)}"
+      )
+    dtype = self.inner_aval.dtype
     if not isinstance(dtype, (jnp.dtype, dtypes.ExtendedDType)):
       dtype = jnp.dtype(dtype)
     return ShapedArrayWithMemorySpace(
-        self.shape, dtype, memory_space=self.memory_space
+        self.inner_aval.shape, dtype, memory_space=self.memory_space
     )
 
   def get_ref_aval(self) -> TransformedRef | state.AbstractRef:
     # TODO(sharadmv): Clean this up. ShapedArrayWithMemorySpace fails when we
     # try to apply JAX ops to it.
-    return state.AbstractRef(
-        jax_core.ShapedArray(self.shape, self.dtype), self.memory_space)
+    return state.AbstractRef(self.inner_aval, self.memory_space)
+
+  @property
+  def dtype(self):
+    return self.inner_aval.dtype
+
+  @property
+  def shape(self):
+    return self.inner_aval.shape
 
 
 class MemorySpace(enum.Enum):
@@ -242,10 +252,13 @@ class MemorySpace(enum.Enum):
   KEY = "key"  # Memory space for PRNG keys.
   HOST = "host"  # Host memory space.
 
-  def __call__(self, shape, dtype):
-    if self == MemorySpace.ANY:
-      return jax.ShapeDtypeStruct(shape, dtype)
-    return MemoryRef(shape, dtype, self)
+  def from_type(self, type: jax_core.AbstractValue) -> MemoryRef:
+    return MemoryRef(type, memory_space=self)
+
+  def __call__(self, shape: tuple[int, ...], dtype: jnp.dtype):
+    # A convenience function for constructing MemoryRef types of ShapedArrays.
+    return self.from_type(jax_core.ShapedArray(shape, dtype))
+
 
   def __str__(self) -> str:
     return self.value

jax/_src/pallas/mosaic/core.py

@@ -23,13 +23,13 @@
 from collections.abc import Mapping
 
 import jax
+import jax.numpy as jnp
 from jax.extend import backend as jex_backend
 from jax._src import core as jax_core
 from jax._src import state
 from jax._src import util
 from jax._src.frozen_dict import FrozenDict
 from jax._src.pallas import core as pallas_core
-import jax.numpy as jnp
 import numpy as np
 
 
@@ -166,9 +166,12 @@ class MemorySpace(enum.Enum):
   def __str__(self) -> str:
     return self.value
 
+  def from_type(self, ty):
+    return pallas_core.MemoryRef(ty, memory_space=self)
+
   def __call__(self, shape: tuple[int, ...], dtype: jnp.dtype):
-    # A convenience function for constructing MemoryRef types.
-    return pallas_core.MemoryRef(shape, dtype, self)
+    # A convenience function for constructing MemoryRef types of ShapedArrays.
+    return self.from_type(jax_core.ShapedArray(shape, dtype))
 
 class dma_semaphore(pallas_core.semaphore_dtype): pass
 
@@ -189,7 +192,8 @@ def __call__(self, shape: tuple[int, ...]):
       dtype = pallas_core.BarrierSemaphore()
     else:
       dtype = pallas_core.Semaphore()
-    return pallas_core.MemoryRef(shape, dtype, MemorySpace.SEMAPHORE)
+    return pallas_core.MemoryRef(jax_core.ShapedArray(shape, dtype),
+                                 MemorySpace.SEMAPHORE)
 
   def get_array_aval(self) -> pallas_core.ShapedArrayWithMemorySpace:
     return self(()).get_array_aval()

jax/_src/pallas/mosaic_gpu/core.py

@@ -150,6 +150,7 @@ def __call__(
       collective: bool | None = None,
       layout: TMEMLayout | None = None,
   ) -> pallas_core.MemoryRef:
+    # TODO(sharadmv): Add HiType constructor support.
     if self == MemorySpace.TMEM:
       if transforms:
         raise ValueError("transforms are not supported for TMEM")
@@ -174,8 +175,9 @@ def __call__(
       if packed is not None or collective is not None or layout is not None:
         raise ValueError("packed, collective and layout arguments are only supported for TMEM.")
       mgpu_layout = None
-    return GPUMemoryRef(shape, dtype, memory_space=self, transforms=transforms,
-                        layout=mgpu_layout, collective=collective)
+    return GPUMemoryRef(jax_core.ShapedArray(shape, dtype), memory_space=self,
+                        transforms=transforms, layout=mgpu_layout,
+                        collective=collective)
 
 
 class SemaphoreType(enum.Enum):
@@ -188,7 +190,8 @@ def __call__(self, shape: tuple[int, ...]):
       dtype = pallas_core.BarrierSemaphore()
     else:
       dtype = pallas_core.Semaphore()
-    return pallas_core.MemoryRef(shape, dtype, MemorySpace.GMEM)
+    return pallas_core.MemoryRef(jax_core.ShapedArray(shape, dtype),
+                                 MemorySpace.GMEM)
 
   def get_array_aval(self) -> jax_core.ShapedArray:
     return self(()).get_array_aval()
@@ -484,8 +487,9 @@ def __init__(self, *refs: _GPUMemoryRefTree):
       object.__setattr__(self, "refs", refs)
       num_bytes = max(map(_ref_group_size, self.refs))
       super().__init__(
-          shape=(num_bytes,),
-          dtype=jnp.int8,
+          inner_aval=jax_core.ShapedArray(
+              (num_bytes,), jnp.int8
+          ),
           memory_space=SMEM,
           transforms=(),
       )
@@ -498,8 +502,10 @@ def __init__(self, *refs: _GPUMemoryRefTree):
             "Some aliased TMEM references are collective and some are not."
         )
       super().__init__(
-          shape=(128, max_cols,),
-          dtype=jnp.int32,
+          inner_aval=jax_core.ShapedArray(
+              shape=(128, max_cols,),
+              dtype=jnp.int32,
+          ),
           memory_space=TMEM,
           transforms=(),
           layout=tcgen05.tmem_default_layout(packing=1),

jax/_src/pallas/mosaic_gpu/lowering.py

@@ -748,11 +748,12 @@ def ref_for_aval(aval: ShapedAbstractValue):
       return gpu_core.WGMMAAccumulatorRef(aval.shape, aval.dtype)
     elif isinstance(aval, gpu_core.AbstractTMEMRef):
       return gpu_core.GPUMemoryRef(
-          aval.shape, aval.dtype, gpu_core.TMEM,
+          jax_core.ShapedArray(aval.shape, aval.dtype), gpu_core.TMEM,
           transforms=(), layout=aval.layout, collective=aval.collective,
       )
     elif isinstance(aval, state_types.AbstractRef):
-      return pallas_core.MemoryRef(aval.shape, aval.dtype, aval.memory_space)
+      return pallas_core.MemoryRef(jax_core.ShapedArray(aval.shape, aval.dtype),
+                                   aval.memory_space)
     else:
       return gpu_core.SMEM(aval.shape, aval.dtype)