[pallas:mgpu] Fix handling of AbstractRef in _commute_transform.

PiperOrigin-RevId: 873912326

Author: chr1sj0nes

jax/_src/pallas/mosaic_gpu/core.py

@@ -978,7 +978,7 @@ def swizzle_elems(self, dtype: jnp.dtype | ir.Type) -> int:
     return (self.swizzle * 8) // mgpu.bitwidth(dtype)
 
   def commute_transpose(
-      self, _: jax_core.ShapedArray, transpose: state_types.TransposeTransform
+      self, _: jax_core.AbstractValue, transpose: state_types.TransposeTransform
   ) -> tuple[state_types.TransposeTransform, UnswizzleRef]:
     perm = transpose.permutation
     if perm[-1] != len(perm) - 1:

jax/_src/pallas/mosaic_gpu/lowering.py

@@ -1454,13 +1454,14 @@ def _commute_transform(
         gpu_core.UnswizzleRef() as t1,
         state_types.ReshapeTransform() as t2,
     ):
-      assert isinstance(aval, jax_core.ShapedArray)
       new_reshape, new_unswizzle = t1.commute_reshape(aval, t2)
       return new_reshape, new_unswizzle
     case (
         gpu_core.UntilingTransform() | gpu_core.UnswizzleRef() as t1,
         gpu_core.TransposeTransform() as t2,
     ):
+      if isinstance(aval, state_types.AbstractRef):
+        aval = aval.inner_aval
       assert isinstance(aval, jax_core.ShapedArray)
       new_reshape, new_unswizzle = t1.commute_transpose(aval, t2)
       return new_reshape, new_unswizzle