[IR] Fix an error when checking for float8_e4m3fnuz type in ir.Tensor (#2078)

justinchuby · web-flow · commit 89dd4544cd31 · 2025-02-27T08:48:01.000-08:00
The float8_e4m3fnuz type was mistaken with float8_e4m3b11fnuz, which is a different type: https://github.com/jax-ml/ml_dtypes#float8_e4m3b11fnuz
diff --git a/onnxscript/ir/_core.py b/onnxscript/ir/_core.py
@@ -199,7 +199,7 @@ def _check_numpy_representation_type(array: np.ndarray, dtype: _enums.DataType)
             )
         if dtype.itemsize == 1 and array.dtype not in (
             np.uint8,
-            ml_dtypes.float8_e4m3b11fnuz,
+            ml_dtypes.float8_e4m3fnuz,
             ml_dtypes.float8_e4m3fn,
             ml_dtypes.float8_e5m2fnuz,
             ml_dtypes.float8_e5m2,

Original file line number	Diff line number	Diff line change
`@@ -199,7 +199,7 @@ def _check_numpy_representation_type(array: np.ndarray, dtype: _enums.DataType)`
`199`	`199`	`)`
`200`	`200`	`if dtype.itemsize == 1 and array.dtype not in (`
`201`	`201`	`np.uint8,`
`202`		`- ml_dtypes.float8_e4m3b11fnuz,`
	`202`	`+ ml_dtypes.float8_e4m3fnuz,`
`203`	`203`	`ml_dtypes.float8_e4m3fn,`
`204`	`204`	`ml_dtypes.float8_e5m2fnuz,`
`205`	`205`	`ml_dtypes.float8_e5m2,`