[Mosaic GPU] Add checks for argument shapes and types

Apparently we never checked it and it's been quite easy to
get this wrong.

PiperOrigin-RevId: 761977944

Author: apaszke

jax/experimental/mosaic/gpu/core.py

@@ -677,14 +677,25 @@ def as_gpu_kernel(
   if launch_ctx.is_device_collective and not supports_cross_device_collectives():
     raise RuntimeError("Kernel is a cross-device collective but no support is available.")
 
-  expected_arg_treedef = jax.tree.structure(in_shape)
+  expected_arg_tys, expected_arg_treedef = jax.tree.flatten(in_shape)
   def _check_args(*args):
     arg_treedef = jax.tree.structure(args)
     if arg_treedef != expected_arg_treedef:
       raise ValueError(
           f"Invalid argument structure: expected {expected_arg_treedef}, got"
           f" {arg_treedef}, ({args=})"
       )
+    for arg, expected_ty in zip(args, expected_arg_tys):
+      if arg.shape != expected_ty.shape:
+        raise ValueError(
+            f"Argument shape mismatch: expected {expected_ty.shape}, got"
+            f" {arg.shape}"
+        )
+      if arg.dtype != expected_ty.dtype:
+        raise ValueError(
+            f"Argument dtype mismatch: expected {expected_ty.dtype}, got"
+            f" {arg.dtype}"
+        )
 
   def bind(*args) -> Any:
     return mosaic_gpu_p.bind(*args, module=module, out_types=out_shape)