onnx_graphsurgeon: Constant with float8_e4m3fnuz values is exported as FLOAT8E4M3FN

[luicarus]

Description

dtype_to_onnx in onnx_graphsurgeon/exporters/onnx_exporter.py currently maps ml_dtypes.float8_e4m3fnuz to FLOAT8E4M3FN instead of FLOAT8E4M3FNUZ:

ml_dtype_to_onnx_name = {
    np.dtype(ml_dtypes.float8_e4m3fn):   "FLOAT8E4M3FN",
    np.dtype(ml_dtypes.float8_e4m3fnuz): "FLOAT8E4M3FN",  # should be FLOAT8E4M3FNUZ
    np.dtype(ml_dtypes.float8_e5m2):     "FLOAT8E5M2",
    np.dtype(ml_dtypes.float8_e5m2fnuz): "FLOAT8E5M2FNUZ",
    ...
}

FLOAT8E4M3FN and FLOAT8E4M3FNUZ are distinct ONNX data types with the same storage width. As a result, the exported raw_data has the expected size and bytes, but the tensor is tagged with the wrong type. Consumers that decode those bytes according to the ONNX type can therefore interpret the values differently.

A minimal reproduction:

import ml_dtypes
import numpy as np
import onnx

from onnx_graphsurgeon.ir.tensor import Constant
from onnx_graphsurgeon.exporters.onnx_exporter import OnnxExporter

values = np.random.random_sample((2, 4)).astype(
    ml_dtypes.float8_e4m3fnuz
)
proto = OnnxExporter.export_tensor_proto(
    Constant("c", values=values)
)

print(proto.data_type)                  # 17: FLOAT8E4M3FN
print(onnx.TensorProto.FLOAT8E4M3FNUZ)  # 18: expected

The issue also breaks import/export round trips once a lazy initializer is materialized. Import already handles FLOAT8E4M3FNUZ correctly, but after accessing .values, re-exporting the tensor changes its type from 18 to 17:

graph = gs.import_onnx(model_with_fnuz_initializer)
tensor = graph.tensors()["w"]

_ = tensor.values  # materialize LazyValues

OnnxExporter.export_tensor_proto(tensor).data_type
# 17 (FLOAT8E4M3FN), originally 18 (FLOAT8E4M3FNUZ)

float8_e5m2fnuz is mapped correctly to FLOAT8E5M2FNUZ in the same table, so the problem appears to be isolated to the float8_e4m3fnuz entry.

Expected behavior

ml_dtypes.float8_e4m3fnuz values should export with:

onnx.TensorProto.FLOAT8E4M3FNUZ

and an imported FLOAT8E4M3FNUZ initializer should preserve its data type after materialization and re-export.

Environment

• onnx_graphsurgeon 0.6.1
• also reproducible on current main
• onnx 1.23.0
• ml_dtypes 0.6.0
• Python 3.13

Provenance

The mapping appears to have been introduced in commit a180e081 (10.8 release changes, 2025-01-31), which added the ml_dtype_to_onnx_name table. The affected entry currently matches the float8_e4m3fn mapping immediately above it.