[pallas:mosaic_gpu] Slightly tweaked the error messages in a few places

PiperOrigin-RevId: 842272894

Author: superbobry

jax/_src/pallas/mosaic_gpu/core.py

@@ -670,7 +670,8 @@ def untransform_reshape(
       self, dtype: jnp.dtype, shape: tuple[int, ...]
   ) -> tuple[tuple[int, ...], state_types.Transform]:
     del dtype
-    raise NotImplementedError("Reshapes don't commute with transposes.")
+    # TODO(slebedev): Support this.
+    raise NotImplementedError("Reshapes don't commute with tiling.")
 
   def untransform_index(
       self, dtype: jnp.dtype | ir.Type, idxs: tuple[Index, ...]

jax/experimental/mosaic/gpu/launch_context.py

@@ -862,7 +862,7 @@ def partition_dim(dim: int, idx: ir.Value, num_chunks: int):
     if max(slice_shape) > 256:
       raise ValueError(
           "Async copies only support copying <=256 elements along each"
-          " dimension"
+          f" dimension, got {tuple(slice_shape)}"
       )
     if (zeroth_bw := slice_shape[-1] * element_bitwidth) % 128 != 0:
       raise ValueError(
@@ -1019,7 +1019,7 @@ def async_copy(
       raise ValueError(
           "Expected the SMEM reference to have the same shape as the"
           f" transformed slice: {tuple(smem_ref_ty.shape)} !="
-          f" {slice_shape[len(squeezed_dims):]}"
+          f" {tuple(slice_shape[len(squeezed_dims):])}"
       )
 
     if implementation == AsyncCopyImplementation.CP_ASYNC:

jax/experimental/mosaic/gpu/wgmma.py

@@ -65,7 +65,8 @@ def value(self) -> fa.FragmentedArray:
   @classmethod
   def zero(cls, m, n, dtype=None, *, is_signed: bool | None = None):
     if m % 64 or n % 8:
-      raise ValueError
+      raise ValueError("WGMMA requires m and n to be multiples of 64 and 8, "
+                       f"got {m} and {n}")
     if is_signed is False:  # pylint: disable=g-bool-id-comparison
       raise TypeError("PTX does not support unsigned WGMMA accumulators")
     f32 = ir.F32Type.get()