fix(onnx-graphsurgeon): export float8_e4m3fnuz with the correct ONNX type

tools/onnx-graphsurgeon/CHANGELOG.md

@@ -8,6 +8,12 @@ Dates are in YYYY-MM-DD format.
 ### Added
 - Added support for `FLOAT4E2M1` tensors and upgraded ONNX version to >= 1.20
 
+### Fixed
+- Fixed a bug where a `Constant` whose values were `ml_dtypes.float8_e4m3fnuz` was exported with the
+  `FLOAT8E4M3FN` ONNX data type instead of `FLOAT8E4M3FNUZ`. The two types share a byte width, so the
+  tensor was silently mislabelled and would be interpreted as the wrong format by anything reading it
+  back — including an import followed by a re-export.
+
 
 ## v0.6.1 (2026-02-17)
 

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",

tools/onnx-graphsurgeon/tests/test_exporters.py

@@ -156,6 +156,111 @@ 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
 
+    @pytest.mark.parametrize(
+        "dtype, expected_onnx_type",
+        [
+            (ml_dtypes.float8_e4m3fn, onnx.TensorProto.FLOAT8E4M3FN),
+            (ml_dtypes.float8_e4m3fnuz, onnx.TensorProto.FLOAT8E4M3FNUZ),
+            (ml_dtypes.float8_e5m2, onnx.TensorProto.FLOAT8E5M2),
+            (ml_dtypes.float8_e5m2fnuz, onnx.TensorProto.FLOAT8E5M2FNUZ),
+        ],
+    )
+    def test_should_export_constant_tensor_with_each_ml_float8_dtype(
+        self, dtype, expected_onnx_type
+    ) -> None:
+        """Test that every `ml_dtypes` float8 variant maps to its own ONNX data type.
+
+        `float8_e4m3fn` and `float8_e4m3fnuz` are distinct ONNX types that share
+        the same byte width, so mapping one to the other silently mislabels the
+        tensor while leaving `raw_data` byte-identical. This asserts the mapping
+        for each variant instead of relying on the bytes alone.
+        """
+        # Precondition.
+        name = "constant_tensor"
+        shape = (3, 224, 224)
+        values = np.random.random_sample(size=shape).astype(dtype)
+
+        tensor = Constant(name=name, values=values)
+
+        # Under test.
+        onnx_tensor = OnnxExporter.export_tensor_proto(tensor)
+
+        # Postcondition.
+        assert onnx_tensor.name == name
+        assert onnx_tensor.data_type == expected_onnx_type
+        assert tuple(onnx_tensor.dims) == shape
+        # The bytes are passed through unchanged regardless of the label.
+        assert bytes(onnx_tensor.raw_data) == values.tobytes()
+
+    def test_should_distinguish_e4m3fn_from_e4m3fnuz(self) -> None:
+        """`float8_e4m3fn` and `float8_e4m3fnuz` must not collapse to the same ONNX type.
+
+        Regression test: the lookup table mapped `float8_e4m3fnuz` to
+        `FLOAT8E4M3FN`, so the two were indistinguishable on export.
+        """
+        shape = (3, 224, 224)
+        values_fn = np.random.random_sample(size=shape).astype(ml_dtypes.float8_e4m3fn)
+        values_fnuz = np.random.random_sample(size=shape).astype(
+            ml_dtypes.float8_e4m3fnuz
+        )
+
+        exported_fn = OnnxExporter.export_tensor_proto(
+            Constant(name="fn", values=values_fn)
+        )
+        exported_fnuz = OnnxExporter.export_tensor_proto(
+            Constant(name="fnuz", values=values_fnuz)
+        )
+
+        assert exported_fn.data_type == onnx.TensorProto.FLOAT8E4M3FN
+        assert exported_fnuz.data_type == onnx.TensorProto.FLOAT8E4M3FNUZ
+        assert exported_fn.data_type != exported_fnuz.data_type
+
+    def test_should_preserve_e4m3fnuz_across_import_and_export_round_trip(
+        self,
+    ) -> None:
+        """A FLOAT8E4M3FNUZ initializer must survive import -> export unchanged.
+
+        Import already handles `FLOAT8E4M3FNUZ` correctly, so a mismatch on
+        export makes the round trip asymmetric: reading `.values` and writing the
+        tensor back out relabels it as `FLOAT8E4M3FN`.
+        """
+        shape = (2, 4)
+        values = np.array(
+            [[0.0, 1.0, 2.0, -2.0], [0.5, 3.0, -3.0, 0.25]], dtype=np.float32
+        ).astype(ml_dtypes.float8_e4m3fnuz)
+
+        onnx_tensor = onnx.helper.make_tensor(
+            name="w",
+            data_type=onnx.TensorProto.FLOAT8E4M3FNUZ,
+            dims=list(shape),
+            vals=values.tobytes(),
+            raw=True,
+        )
+        model = onnx.helper.make_model(
+            onnx.helper.make_graph(
+                nodes=[onnx.helper.make_node("Identity", ["w"], ["y"])],
+                name="graph",
+                inputs=[],
+                outputs=[
+                    onnx.helper.make_tensor_value_info(
+                        "y", onnx.TensorProto.FLOAT8E4M3FNUZ, list(shape)
+                    )
+                ],
+                initializer=[onnx_tensor],
+            ),
+            opset_imports=[onnx.helper.make_opsetid("", 20)],
+        )
+
+        # Import, force LazyValues to load (as any weight-editing pass would),
+        # then export again.
+        graph = OnnxImporter.import_graph(model.graph)
+        tensor = graph.tensors()["w"]
+        _ = tensor.values
+
+        re_exported = OnnxExporter.export_tensor_proto(tensor)
+
+        assert re_exported.data_type == onnx.TensorProto.FLOAT8E4M3FNUZ
+
     def test_should_export_constant_tensor_with_ml_dtype_raise_error_when_onnx_dtype_not_supported(
         self,
     ) -> None: