[mgpu] Added a missed case for debug_print types and raise a proper error if a type is unexpected.

PiperOrigin-RevId: 701003002

Author: cperivol

jax/experimental/mosaic/gpu/utils.py

@@ -108,21 +108,18 @@ def c(val: int | float, ty):
   return arith.constant(ty, attr)
 
 def _debug_scalar_ty_format(arg):
-  ty_format = None
   if ir.IndexType.isinstance(arg.type):
-    return "%llu"
+    return "%llu", arg
   if ir.IntegerType.isinstance(arg.type):
-    width = ir.IntegerType(arg.type).width
-    ty_format = "%llu"
-    if width < 64:
+    if ir.IntegerType(arg.type).width < 64:
       arg = arith.extui(ir.IntegerType.get_signless(64), arg)
+    return "%llu", arg
   if ir.F32Type.isinstance(arg.type):
-    ty_format = "%f"
+    return "%f", arg
   if ir.F16Type.isinstance(arg.type):
-    ty_format = "%f"
     arg = arith.extf(ir.F32Type.get(), arg)
-
-  return ty_format, arg
+    return "%f", arg
+  raise NotImplementedError(f"Can't print the type {arg.type}")
 
 def debug_print(fmt, *args, uniform=True):
   type_formats = []