From fbfcfc164b9bdc6b2d59048e8fdb3d6a17363521 Mon Sep 17 00:00:00 2001 From: rootkiller6788 Date: Mon, 31 Aug 2026 23:12:15 +0800 Subject: [PATCH 1/2] fix: export fp8e4m3fnuz tensors with the correct onnx type dtype_to_onnx mapped ml_dtypes float8_e4m3fnuz to FLOAT8E4M3FN instead of FLOAT8E4M3FNUZ (copy-paste of the fn entry). A constant with fnuz values would get silently relabeled as fn on export. --- .../onnx_graphsurgeon/exporters/onnx_exporter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/onnx-graphsurgeon/onnx_graphsurgeon/exporters/onnx_exporter.py b/tools/onnx-graphsurgeon/onnx_graphsurgeon/exporters/onnx_exporter.py index fa5fa0de..35e4982d 100644 --- a/tools/onnx-graphsurgeon/onnx_graphsurgeon/exporters/onnx_exporter.py +++ b/tools/onnx-graphsurgeon/onnx_graphsurgeon/exporters/onnx_exporter.py @@ -50,7 +50,7 @@ def dtype_to_onnx(dtype: Union[np.dtype, "onnx.TensorProto.DataType"]) -> int: ml_dtype_to_onnx_name = { np.dtype(ml_dtypes.bfloat16): "BFLOAT16", np.dtype(ml_dtypes.float8_e4m3fn): "FLOAT8E4M3FN", - np.dtype(ml_dtypes.float8_e4m3fnuz): "FLOAT8E4M3FN", + np.dtype(ml_dtypes.float8_e4m3fnuz): "FLOAT8E4M3FNUZ", np.dtype(ml_dtypes.float8_e5m2): "FLOAT8E5M2", np.dtype(ml_dtypes.float8_e5m2fnuz): "FLOAT8E5M2FNUZ", np.dtype(ml_dtypes.uint4): "UINT4", From 60981d9a03eb0d87768cb8207c45d924dee2e8b4 Mon Sep 17 00:00:00 2001 From: rootkiller6788 Date: Mon, 31 Aug 2026 23:12:26 +0800 Subject: [PATCH 2/2] test: add fp8e4m3fnuz constant export regression test --- .../onnx-graphsurgeon/tests/test_exporters.py | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/tools/onnx-graphsurgeon/tests/test_exporters.py b/tools/onnx-graphsurgeon/tests/test_exporters.py index a5eaf6c2..d46f6538 100644 --- a/tools/onnx-graphsurgeon/tests/test_exporters.py +++ b/tools/onnx-graphsurgeon/tests/test_exporters.py @@ -156,6 +156,27 @@ def test_should_export_constant_tensor_with_ml_dtype(self) -> None: assert onnx_tensor.data_type == onnx.TensorProto.FLOAT8E4M3FN assert tuple(onnx_tensor.dims) == shape + def test_should_export_constant_tensor_with_ml_dtype_fnuz(self) -> None: + """Test that `export_tensor_proto` exports a Constant with ml_dtypes.float8_e4m3fnuz values + using the correct ONNX data type.""" + name = "constant_tensor" + shape = (3, 224, 224) + values = np.random.random_sample(size=shape).astype(ml_dtypes.float8_e4m3fnuz) + + tensor = Constant(name=name, values=values) + + onnx_tensor = OnnxExporter.export_tensor_proto(tensor) + + assert onnx_tensor.name == name + assert np.all( + self._bytes_to_np_array( + onnx_tensor.raw_data, onnx_tensor.dims, ml_dtypes.float8_e4m3fnuz + ) + == tensor.values + ) + assert onnx_tensor.data_type == onnx.TensorProto.FLOAT8E4M3FNUZ + assert tuple(onnx_tensor.dims) == shape + def test_should_export_constant_tensor_with_ml_dtype_raise_error_when_onnx_dtype_not_supported( self, ) -> None: