Added a missing branch to mgpu.FragmentedArray.astype

Previously, an unsupported cast produced a `NameError` instead. PiperOrigin-RevId: 676010161
keshavb96 · Sep 18, 2024 · 442e863 · 442e863
1 parent 6236b8f
commit 442e863
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/jax/experimental/mosaic/gpu/fragmented_array.py b/jax/experimental/mosaic/gpu/fragmented_array.py
@@ -494,6 +494,8 @@ def astype(self, new_dtype: ir.Type):
       convert = arith.sitofp
     elif from_float and to_integer:
       convert = arith.fptosi
+    else:
+      raise NotImplementedError(f"Unsupported conversion {cur_dtype} -> {new_dtype}")
     new_registers = np.empty_like(self.registers)
     match self.layout:
       case WGMMAFragLayout():