Tiny usability tweak to allow user to pass a list instead of a tuple, e.g. pltpu.VMEM([8], jnp.int32).

PiperOrigin-RevId: 836134234

Author: brianwa84

jax/_src/pallas/mosaic/core.py

@@ -183,9 +183,9 @@ def __str__(self) -> str:
   def from_type(self, ty):
     return pallas_core.MemoryRef(ty, memory_space=self)
 
-  def __call__(self, shape: tuple[int, ...], dtype: jnp.dtype):
+  def __call__(self, shape: Sequence[int], dtype: jnp.dtype):
     # A convenience function for constructing MemoryRef types of ShapedArrays.
-    return self.from_type(jax_core.ShapedArray(shape, dtype))
+    return self.from_type(jax_core.ShapedArray(tuple(shape), dtype))
 
 class dma_semaphore(pallas_core.semaphore_dtype): pass
 

jax/_src/pallas/mosaic_gpu/core.py

@@ -142,14 +142,15 @@ def __str__(self) -> str:
 
   def __call__(
       self,
-      shape: tuple[int, ...],
+      shape: Sequence[int],
       dtype: jnp.dtype,
       *,
       transforms: Sequence[MemoryRefTransform] = (),
       packed: bool | None = None,
       collective: bool | None = None,
       layout: TMEMLayout | None = None,
   ) -> pallas_core.MemoryRef:
+    shape = tuple(shape)
     # TODO(sharadmv): Add HiType constructor support.
     if self == MemorySpace.TMEM:
       if transforms: