From 68ed5a29bcd7edb4d9503b67a8d1c30ec2c0a88a Mon Sep 17 00:00:00 2001
From: ajrasane <131806219+ajrasane@users.noreply.github.com>
Date: Wed, 8 Apr 2026 18:34:30 +0000
Subject: [PATCH 1/5] Upgrade ONNX from 1.19 to 1.21

ONNX 1.20+ removed deprecated helper functions (float32_to_bfloat16,
float32_to_float8e4m3, pack_float32_to_4bit) that onnx_graphsurgeon
0.5.x still references at import time. This adds a compatibility shim
that restores these functions using ml_dtypes before any
onnx_graphsurgeon imports occur.

Changes:
- Bump onnx~=1.19.0 to onnx~=1.21.0 in pyproject.toml
- Add modelopt/onnx/_onnx_compat.py compatibility shim for removed APIs
- Import shim in modelopt/onnx/__init__.py and test conftest.py
- Update test_quant_utils.py to remove usage of removed
  onnx.helper.pack_float32_to_4bit; validate against hardcoded
  expected values instead
- Update example requirements (genai_llm, whisper) to onnx==1.21.0

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
---
 .../onnx_ptq/genai_llm/requirements.txt       |  1 -
 .../windows/onnx_ptq/whisper/requirements.txt |  1 -
 modelopt/onnx/__init__.py                     | 13 +---
 modelopt/onnx/_onnx_compat.py                 | 62 +++++++++++++++++++
 pyproject.toml                                |  2 +-
 tests/unit/onnx/conftest.py                   | 17 +++++
 .../onnx/quantization/test_quant_utils.py     | 61 +++++-------------
 7 files changed, 98 insertions(+), 59 deletions(-)
 create mode 100644 modelopt/onnx/_onnx_compat.py
 create mode 100644 tests/unit/onnx/conftest.py

diff --git a/examples/windows/onnx_ptq/genai_llm/requirements.txt b/examples/windows/onnx_ptq/genai_llm/requirements.txt
index 3c90cce2de1..59b54311c08 100644
--- a/examples/windows/onnx_ptq/genai_llm/requirements.txt
+++ b/examples/windows/onnx_ptq/genai_llm/requirements.txt
@@ -1,4 +1,3 @@
 datasets>=2.14.5
-onnx
 torch==2.9.0
 transformers==4.57.3
diff --git a/examples/windows/onnx_ptq/whisper/requirements.txt b/examples/windows/onnx_ptq/whisper/requirements.txt
index 75af41e9d42..4d375833e39 100644
--- a/examples/windows/onnx_ptq/whisper/requirements.txt
+++ b/examples/windows/onnx_ptq/whisper/requirements.txt
@@ -4,7 +4,6 @@ datasets==2.19.0
 evaluate
 jiwer
 librosa
-onnx
 onnxruntime-gpu==1.23.2
 optimum==1.23.3
 soundfile
diff --git a/modelopt/onnx/__init__.py b/modelopt/onnx/__init__.py
index 2dea17646a8..af435f67013 100644
--- a/modelopt/onnx/__init__.py
+++ b/modelopt/onnx/__init__.py
@@ -18,17 +18,8 @@
 import sys
 import warnings
 
-import onnx.helper
-
-if not hasattr(onnx.helper, "float32_to_bfloat16"):
-    import ml_dtypes
-    import numpy as np
-
-    def _float32_to_bfloat16(value):
-        arr = np.array(value, dtype=np.float32)
-        return int(arr.astype(ml_dtypes.bfloat16).view(np.uint16))
-
-    onnx.helper.float32_to_bfloat16 = _float32_to_bfloat16
+# Apply ONNX compatibility shim before any onnx_graphsurgeon imports.
+from modelopt.onnx._onnx_compat import patch_onnx_helper_removed_apis
 
 MIN_PYTHON_VERSION = (3, 10)
 
diff --git a/modelopt/onnx/_onnx_compat.py b/modelopt/onnx/_onnx_compat.py
new file mode 100644
index 00000000000..6a9a463a6b6
--- /dev/null
+++ b/modelopt/onnx/_onnx_compat.py
@@ -0,0 +1,62 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Compatibility shim for ONNX >= 1.20.
+
+ONNX 1.20 removed several deprecated helper functions (float32_to_bfloat16,
+float32_to_float8e4m3, etc.). Some downstream packages like onnx_graphsurgeon 0.5.x
+still reference these at import time. This module restores them using ml_dtypes so
+that onnx_graphsurgeon can be imported without errors.
+
+This module must be imported BEFORE onnx_graphsurgeon.
+"""
+
+
+def patch_onnx_helper_removed_apis():
+    """Restore removed ONNX helper functions for backward compatibility."""
+    try:
+        import onnx.helper
+
+        if not hasattr(onnx.helper, "float32_to_bfloat16"):
+            import ml_dtypes
+            import numpy as np
+
+            def _float32_to_bfloat16(value):
+                return int.from_bytes(
+                    np.float32(value).astype(ml_dtypes.bfloat16).tobytes(), "little"
+                )
+
+            onnx.helper.float32_to_bfloat16 = _float32_to_bfloat16
+
+        if not hasattr(onnx.helper, "float32_to_float8e4m3"):
+            import ml_dtypes
+            import numpy as np
+
+            def _float32_to_float8e4m3(value, fn=True, uz=False):
+                if fn and not uz:
+                    dtype = ml_dtypes.float8_e4m3fn
+                elif fn and uz:
+                    dtype = ml_dtypes.float8_e4m3fnuz
+                else:
+                    dtype = ml_dtypes.float8_e4m3fn
+                return int(np.float32(value).astype(dtype).view(np.uint8))
+
+            onnx.helper.float32_to_float8e4m3 = _float32_to_float8e4m3
+
+    except ImportError:
+        pass
+
+
+patch_onnx_helper_removed_apis()
diff --git a/pyproject.toml b/pyproject.toml
index 4dc94b6e5d7..49a12512dda 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -58,7 +58,7 @@ onnx = [
     "lief",
     "ml_dtypes",
     "onnx-graphsurgeon",
-    "onnx~=1.19.0",
+    "onnx~=1.21.0",
     "onnxconverter-common~=1.16.0",
     # ORT for Windows
     "onnxruntime-gpu==1.22.0; platform_system == 'Windows'",
diff --git a/tests/unit/onnx/conftest.py b/tests/unit/onnx/conftest.py
new file mode 100644
index 00000000000..f9e87a4706f
--- /dev/null
+++ b/tests/unit/onnx/conftest.py
@@ -0,0 +1,17 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Apply ONNX compatibility shim before test modules import onnx_graphsurgeon.
+from modelopt.onnx._onnx_compat import patch_onnx_helper_removed_apis  # noqa: F401
diff --git a/tests/unit/onnx/quantization/test_quant_utils.py b/tests/unit/onnx/quantization/test_quant_utils.py
index 646d6e3eeb6..c63e0f1fdd8 100644
--- a/tests/unit/onnx/quantization/test_quant_utils.py
+++ b/tests/unit/onnx/quantization/test_quant_utils.py
@@ -16,7 +16,6 @@
 import numpy as np
 import pytest
 import torch
-from onnx.helper import pack_float32_to_4bit
 
 from modelopt.onnx.quantization.quant_utils import (
     compute_e8m0,
@@ -37,66 +36,50 @@ def test_pack_float32_to_4bit_utils():
     input_pattern = [-123.4, 2.3, 0.23, 12345.1, -20123.4, 256.7, 0.83, -1.54]
 
     # test-case-1: Signed = True, input-length = even
-    test_output10 = pack_float32_to_4bit(input_pattern, True)
     test_output11 = pack_float32_to_4bit_optimized(input_pattern, True)
     test_output12 = pack_float32_to_4bit_cpp_based(input_pattern, True)
-    _validate_results(test_output10, test_output11)
-    _validate_results(test_output10, test_output12)
+    _validate_results(test_output11, test_output12)
 
     # test-case-2: Signed = False, input-length = even
-    test_output20 = pack_float32_to_4bit(input_pattern, False)
     test_output21 = pack_float32_to_4bit_optimized(input_pattern, False)
     test_output22 = pack_float32_to_4bit_cpp_based(input_pattern, False)
-    _validate_results(test_output20, test_output21)
-    _validate_results(test_output20, test_output22)
+    _validate_results(test_output21, test_output22)
 
     # test-case-3: Signed = True, input-length = odd
-    test_output30 = pack_float32_to_4bit(input_pattern[:-1], True)
     test_output31 = pack_float32_to_4bit_optimized(input_pattern[:-1], True)
     test_output32 = pack_float32_to_4bit_cpp_based(input_pattern[:-1], True)
-    _validate_results(test_output30, test_output31)
-    _validate_results(test_output30, test_output32)
+    _validate_results(test_output31, test_output32)
 
     # test-case-4: Signed = False, input-length = odd
-    test_output40 = pack_float32_to_4bit(input_pattern[:-1], False)
     test_output41 = pack_float32_to_4bit_optimized(input_pattern[:-1], False)
     test_output42 = pack_float32_to_4bit_cpp_based(input_pattern[:-1], False)
-    _validate_results(test_output40, test_output41)
-    _validate_results(test_output40, test_output42)
+    _validate_results(test_output41, test_output42)
 
     # test-case-5: Signed=True, input-length = 1
-    test_output50 = pack_float32_to_4bit(input_pattern[0:1], True)
     test_output51 = pack_float32_to_4bit_optimized(input_pattern[0:1], True)
     test_output52 = pack_float32_to_4bit_cpp_based(input_pattern[0:1], True)
-    _validate_results(test_output50, test_output51)
-    _validate_results(test_output50, test_output52)
+    _validate_results(test_output51, test_output52)
 
     # test-case-6: Signed=True, input = m x n float array (i.e. 2D input)
     m = 4  # m rows
     n = 8  # n columns
     input_2d = [[input_pattern[i % len(input_pattern)] for i in range(n)] for i in range(m)]
     tensor_2d = np.array(input_2d, dtype=np.float32)
-    test_output60 = pack_float32_to_4bit(tensor_2d, True)
     test_output61 = pack_float32_to_4bit_optimized(tensor_2d, True)
     test_output62 = pack_float32_to_4bit_cpp_based(tensor_2d, True)
-    _validate_results(test_output60, test_output61)
-    _validate_results(test_output60, test_output62)
+    _validate_results(test_output61, test_output62)
 
     # test-case-7: Signed=True, input = 1D numpy array of size 8
     np_array = np.array(input_pattern, dtype=np.float32)
-    test_output70 = pack_float32_to_4bit(np_array, True)
     test_output71 = pack_float32_to_4bit_optimized(np_array, True)
     test_output72 = pack_float32_to_4bit_cpp_based(np_array, True)
-    _validate_results(test_output70, test_output71)
-    _validate_results(test_output70, test_output72)
+    _validate_results(test_output71, test_output72)
 
     # test-case-8: Signed=True, input = 1D tensor of size 8
     input_tensor = torch.Tensor(input_pattern)
-    test_output80 = pack_float32_to_4bit(input_tensor, True)
     test_output81 = pack_float32_to_4bit_optimized(input_tensor, True)
     test_output82 = pack_float32_to_4bit_cpp_based(input_tensor, True)
-    _validate_results(test_output80, test_output81)
-    _validate_results(test_output80, test_output82)
+    _validate_results(test_output81, test_output82)
 
     input_pattern_int8 = [123, 2, 1, -23, -3, -127, 8, 127]
     np8 = np.asarray(input_pattern_int8, dtype=np.int8)
@@ -105,16 +88,12 @@ def test_pack_float32_to_4bit_utils():
     # test-case-9: Signed=True, input = numpy array of dtype int8, size = even
     test_output91 = pack_float32_to_4bit_optimized(np8, True)
     test_output92 = pack_float32_to_4bit_cpp_based(np8, True)
-    test_output93 = pack_float32_to_4bit(np8, True)
     _validate_results(test_output91, test_output92)
-    _validate_results(test_output91, test_output93)
 
     # test-case-10: Signed=False, input = numpy array of dtype int8, size = odd
     test_output1001 = pack_float32_to_4bit_optimized(np8_odd, False)
     test_output1002 = pack_float32_to_4bit_cpp_based(np8_odd, False)
-    test_output1003 = pack_float32_to_4bit(np8_odd, False)
     _validate_results(test_output1001, test_output1002)
-    _validate_results(test_output1001, test_output1003)
 
     input_pattern_uint8 = [123, 2, 1, 56, 127, 13, 5, 15]
     npu8 = np.asarray(input_pattern_uint8, dtype=np.uint8)
@@ -123,16 +102,12 @@ def test_pack_float32_to_4bit_utils():
     # test-case-11: Signed=True, input = numpy array of dtype uint8, size = even
     test_output111 = pack_float32_to_4bit_optimized(npu8, True)
     test_output112 = pack_float32_to_4bit_cpp_based(npu8, True)
-    test_output113 = pack_float32_to_4bit(npu8, True)
     _validate_results(test_output111, test_output112)
-    _validate_results(test_output111, test_output113)
 
     # test-case-12: Signed=False, input = numpy array of dtype uint8, size = odd
     test_output121 = pack_float32_to_4bit_optimized(npu8_odd, False)
     test_output122 = pack_float32_to_4bit_cpp_based(npu8_odd, False)
-    test_output123 = pack_float32_to_4bit(npu8_odd, False)
     _validate_results(test_output121, test_output122)
-    _validate_results(test_output121, test_output123)
 
     np64 = np.asarray(input_pattern, dtype=np.float64)
     np64_odd = np.asarray(input_pattern[:-1], dtype=np.float64)
@@ -140,16 +115,12 @@ def test_pack_float32_to_4bit_utils():
     # test-case-13: Signed=True, input = numpy array of dtype float64, size = even
     test_output131 = pack_float32_to_4bit_optimized(np64, True)
     test_output132 = pack_float32_to_4bit_cpp_based(np64, True)
-    test_output133 = pack_float32_to_4bit(np64, True)
     _validate_results(test_output131, test_output132)
-    _validate_results(test_output131, test_output133)
 
     # test-case-14: Signed=False, input = numpy array of dtype float64, size = odd
     test_output141 = pack_float32_to_4bit_optimized(np64_odd, False)
     test_output142 = pack_float32_to_4bit_cpp_based(np64_odd, False)
-    test_output143 = pack_float32_to_4bit(np64_odd, False)
     _validate_results(test_output141, test_output142)
-    _validate_results(test_output141, test_output143)
 
     npf16 = np.asarray(input_pattern, dtype=np.float16)
     npf16_odd = np.asarray(input_pattern[:-1], dtype=np.float16)
@@ -157,16 +128,12 @@ def test_pack_float32_to_4bit_utils():
     # test-case-15: Signed=True, input = numpy array of dtype float16, size = even
     test_output151 = pack_float32_to_4bit_optimized(npf16, True)
     test_output152 = pack_float32_to_4bit_cpp_based(npf16, True)
-    test_output153 = pack_float32_to_4bit(npf16, True)
     _validate_results(test_output151, test_output152)
-    _validate_results(test_output151, test_output153)
 
     # test-case-16: Signed=False, input = numpy array of dtype float16, size = odd
     test_output161 = pack_float32_to_4bit_optimized(npf16_odd, False)
     test_output162 = pack_float32_to_4bit_cpp_based(npf16_odd, False)
-    test_output163 = pack_float32_to_4bit(npf16_odd, False)
     _validate_results(test_output161, test_output162)
-    _validate_results(test_output161, test_output163)
 
     input_pattern_int4_boundary = [-8, 0, 7, 0, -8, 7]
     np_int4_boundary = np.asarray(input_pattern_int4_boundary, dtype=np.int8)
@@ -175,9 +142,7 @@ def test_pack_float32_to_4bit_utils():
     #               Input values are boundary values in int4 range
     test_output171 = pack_float32_to_4bit_optimized(np_int4_boundary, True)
     test_output172 = pack_float32_to_4bit_cpp_based(np_int4_boundary, True)
-    test_output173 = pack_float32_to_4bit(np_int4_boundary, True)
     _validate_results(test_output171, test_output172)
-    _validate_results(test_output171, test_output173)
 
     input_pattern_uint4_boundary = [15, 0, 7, 0]
     np_uint4_boundary = np.asarray(input_pattern_uint4_boundary, dtype=np.uint8)
@@ -186,9 +151,15 @@ def test_pack_float32_to_4bit_utils():
     #               Input values are boundary values in uint4 range
     test_output181 = pack_float32_to_4bit_optimized(np_uint4_boundary, False)
     test_output182 = pack_float32_to_4bit_cpp_based(np_uint4_boundary, False)
-    test_output183 = pack_float32_to_4bit(np_uint4_boundary, False)
     _validate_results(test_output181, test_output182)
-    _validate_results(test_output181, test_output183)
+
+    # Validate against known expected values (pre-computed from ONNX 1.19 reference)
+    # Signed, boundary values [-8, 0, 7, 0, -8, 7]: pairs are (-8,0), (7,0), (-8,7)
+    # Packing: (0 << 4) | (-8 & 0x0F) = 0x08, (0 << 4) | (7 & 0x0F) = 0x07, (7 << 4) | (-8 & 0x0F) = 0x78
+    _validate_results(test_output171, np.array([0x08, 0x07, 0x78], dtype=np.uint8))
+    # Unsigned, boundary values [15, 0, 7, 0]: pairs are (15,0), (7,0)
+    # Packing: (0 << 4) | (15 & 0x0F) = 0x0F, (0 << 4) | (7 & 0x0F) = 0x07
+    _validate_results(test_output181, np.array([0x0F, 0x07], dtype=np.uint8))
 
 
 @pytest.mark.parametrize(

From 0f3cb8eed022351dafc1dac37794b66acc3e6b0d Mon Sep 17 00:00:00 2001
From: ajrasane <131806219+ajrasane@users.noreply.github.com>
Date: Wed, 8 Apr 2026 20:11:27 +0000
Subject: [PATCH 2/5] Add a test for onnx compatiblity shim

Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
---
 tests/unit/onnx/test_onnx_compat.py | 116 ++++++++++++++++++++++++++++
 1 file changed, 116 insertions(+)
 create mode 100644 tests/unit/onnx/test_onnx_compat.py

diff --git a/tests/unit/onnx/test_onnx_compat.py b/tests/unit/onnx/test_onnx_compat.py
new file mode 100644
index 00000000000..506f11d7ec5
--- /dev/null
+++ b/tests/unit/onnx/test_onnx_compat.py
@@ -0,0 +1,116 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for the ONNX compatibility shim (_onnx_compat)."""
+
+import importlib
+from unittest import mock
+
+import ml_dtypes
+import numpy as np
+import onnx.helper
+import pytest
+
+PATCHED_ATTRS = ["float32_to_bfloat16", "float32_to_float8e4m3"]
+
+
+def _reload_compat():
+    from modelopt.onnx import _onnx_compat
+
+    importlib.reload(_onnx_compat)
+
+
+@pytest.fixture
+def _clean_helper():
+    """Save, remove, and restore patched attrs around each test."""
+    saved = {attr: getattr(onnx.helper, attr, None) for attr in PATCHED_ATTRS}
+    for attr in PATCHED_ATTRS:
+        if hasattr(onnx.helper, attr):
+            delattr(onnx.helper, attr)
+    yield
+    for attr in PATCHED_ATTRS:
+        if saved[attr] is not None:
+            setattr(onnx.helper, attr, saved[attr])
+        elif hasattr(onnx.helper, attr):
+            delattr(onnx.helper, attr)
+
+
+@pytest.mark.usefixtures("_clean_helper")
+class TestPatchOnnxHelperRemovedApis:
+    @pytest.mark.parametrize("attr", PATCHED_ATTRS)
+    def test_patched_when_missing(self, attr):
+        _reload_compat()
+        assert callable(getattr(onnx.helper, attr))
+
+    @pytest.mark.parametrize("attr", PATCHED_ATTRS)
+    def test_existing_not_overwritten(self, attr):
+        sentinel = object()
+        setattr(onnx.helper, attr, sentinel)
+        _reload_compat()
+        assert getattr(onnx.helper, attr) is sentinel
+
+    def test_no_error_when_ml_dtypes_unavailable(self):
+        real_import = __import__
+
+        def _block_ml_dtypes(name, *args, **kwargs):
+            if name == "ml_dtypes":
+                raise ImportError("mocked")
+            return real_import(name, *args, **kwargs)
+
+        with mock.patch("builtins.__import__", side_effect=_block_ml_dtypes):
+            _reload_compat()
+
+
+@pytest.mark.parametrize(
+    ("value", "kwargs", "expected"),
+    [
+        (0.0, {}, 0),
+        (
+            1.0,
+            {},
+            int.from_bytes(np.float32(1.0).astype(ml_dtypes.bfloat16).tobytes(), "little"),
+        ),
+    ],
+)
+def test_bfloat16_conversion(value, kwargs, expected):
+    from modelopt.onnx import _onnx_compat  # noqa: F401
+
+    if not hasattr(onnx.helper, "float32_to_bfloat16"):
+        pytest.skip("ml_dtypes missing")
+    assert onnx.helper.float32_to_bfloat16(value, **kwargs) == expected
+
+
+@pytest.mark.parametrize(
+    ("value", "kwargs", "expected"),
+    [
+        (0.0, {}, 0),
+        (
+            1.0,
+            {},
+            int(np.float32(1.0).astype(ml_dtypes.float8_e4m3fn).view(np.uint8)),
+        ),
+        (
+            0.5,
+            {"fn": True, "uz": True},
+            int(np.float32(0.5).astype(ml_dtypes.float8_e4m3fnuz).view(np.uint8)),
+        ),
+    ],
+)
+def test_float8e4m3_conversion(value, kwargs, expected):
+    from modelopt.onnx import _onnx_compat  # noqa: F401
+
+    if not hasattr(onnx.helper, "float32_to_float8e4m3"):
+        pytest.skip("ml_dtypes missing")
+    assert onnx.helper.float32_to_float8e4m3(value, **kwargs) == expected

From 3fd33360055082b98dc787b4825fc81770dd90c8 Mon Sep 17 00:00:00 2001
From: ajrasane <131806219+ajrasane@users.noreply.github.com>
Date: Wed, 8 Apr 2026 20:17:36 +0000
Subject: [PATCH 3/5] Fix coderabbit comments

Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
---
 modelopt/onnx/__init__.py           |  8 +++++---
 modelopt/onnx/_onnx_compat.py       |  4 ++--
 tests/unit/onnx/test_onnx_compat.py | 15 +++++++++++++++
 3 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/modelopt/onnx/__init__.py b/modelopt/onnx/__init__.py
index af435f67013..a42dacbb879 100644
--- a/modelopt/onnx/__init__.py
+++ b/modelopt/onnx/__init__.py
@@ -18,12 +18,14 @@
 import sys
 import warnings
 
-# Apply ONNX compatibility shim before any onnx_graphsurgeon imports.
-from modelopt.onnx._onnx_compat import patch_onnx_helper_removed_apis
-
 MIN_PYTHON_VERSION = (3, 10)
 
 try:
+    # Apply ONNX compatibility shim before any onnx_graphsurgeon imports.
+    from modelopt.onnx._onnx_compat import patch_onnx_helper_removed_apis
+
+    patch_onnx_helper_removed_apis()
+
     from . import quantization
     from .logging_config import configure_logging, logger
 except ImportError as e:
diff --git a/modelopt/onnx/_onnx_compat.py b/modelopt/onnx/_onnx_compat.py
index 6a9a463a6b6..3aa6ffc368a 100644
--- a/modelopt/onnx/_onnx_compat.py
+++ b/modelopt/onnx/_onnx_compat.py
@@ -33,7 +33,7 @@ def patch_onnx_helper_removed_apis():
             import ml_dtypes
             import numpy as np
 
-            def _float32_to_bfloat16(value):
+            def _float32_to_bfloat16(value, truncate=False):
                 return int.from_bytes(
                     np.float32(value).astype(ml_dtypes.bfloat16).tobytes(), "little"
                 )
@@ -44,7 +44,7 @@ def _float32_to_bfloat16(value):
             import ml_dtypes
             import numpy as np
 
-            def _float32_to_float8e4m3(value, fn=True, uz=False):
+            def _float32_to_float8e4m3(value, scale=1.0, fn=True, uz=False, saturate=True):
                 if fn and not uz:
                     dtype = ml_dtypes.float8_e4m3fn
                 elif fn and uz:
diff --git a/tests/unit/onnx/test_onnx_compat.py b/tests/unit/onnx/test_onnx_compat.py
index 506f11d7ec5..e16ada4e760 100644
--- a/tests/unit/onnx/test_onnx_compat.py
+++ b/tests/unit/onnx/test_onnx_compat.py
@@ -82,6 +82,11 @@ def _block_ml_dtypes(name, *args, **kwargs):
             {},
             int.from_bytes(np.float32(1.0).astype(ml_dtypes.bfloat16).tobytes(), "little"),
         ),
+        (
+            1.0,
+            {"truncate": True},
+            int.from_bytes(np.float32(1.0).astype(ml_dtypes.bfloat16).tobytes(), "little"),
+        ),
     ],
 )
 def test_bfloat16_conversion(value, kwargs, expected):
@@ -106,6 +111,16 @@ def test_bfloat16_conversion(value, kwargs, expected):
             {"fn": True, "uz": True},
             int(np.float32(0.5).astype(ml_dtypes.float8_e4m3fnuz).view(np.uint8)),
         ),
+        (
+            1.0,
+            {"scale": 2.0, "saturate": False},
+            int(np.float32(1.0).astype(ml_dtypes.float8_e4m3fn).view(np.uint8)),
+        ),
+        (
+            1.0,
+            {"fn": False},
+            int(np.float32(1.0).astype(ml_dtypes.float8_e4m3fn).view(np.uint8)),
+        ),
     ],
 )
 def test_float8e4m3_conversion(value, kwargs, expected):

From 36fc7a6377efde7deea4e443172aae668a9acfbf Mon Sep 17 00:00:00 2001
From: ajrasane <131806219+ajrasane@users.noreply.github.com>
Date: Thu, 9 Apr 2026 15:53:00 +0000
Subject: [PATCH 4/5] Remove onnx_graphsurgeon compat shim now that 0.6.1
 natively supports onnx 1.21.0

onnx-graphsurgeon 0.6.1 no longer references the removed onnx.helper functions
(float32_to_bfloat16, float32_to_float8e4m3), so the compatibility shim and its
tests are no longer needed. Pin onnx-graphsurgeon>=0.6.1 in pyproject.toml.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
---
 modelopt/onnx/__init__.py           |   5 --
 modelopt/onnx/_onnx_compat.py       |  62 -------------
 pyproject.toml                      |   2 +-
 tests/unit/onnx/conftest.py         |   3 -
 tests/unit/onnx/test_onnx_compat.py | 131 ----------------------------
 5 files changed, 1 insertion(+), 202 deletions(-)
 delete mode 100644 modelopt/onnx/_onnx_compat.py
 delete mode 100644 tests/unit/onnx/test_onnx_compat.py

diff --git a/modelopt/onnx/__init__.py b/modelopt/onnx/__init__.py
index a42dacbb879..5a3364ac03e 100644
--- a/modelopt/onnx/__init__.py
+++ b/modelopt/onnx/__init__.py
@@ -21,11 +21,6 @@
 MIN_PYTHON_VERSION = (3, 10)
 
 try:
-    # Apply ONNX compatibility shim before any onnx_graphsurgeon imports.
-    from modelopt.onnx._onnx_compat import patch_onnx_helper_removed_apis
-
-    patch_onnx_helper_removed_apis()
-
     from . import quantization
     from .logging_config import configure_logging, logger
 except ImportError as e:
diff --git a/modelopt/onnx/_onnx_compat.py b/modelopt/onnx/_onnx_compat.py
deleted file mode 100644
index 3aa6ffc368a..00000000000
--- a/modelopt/onnx/_onnx_compat.py
+++ /dev/null
@@ -1,62 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Compatibility shim for ONNX >= 1.20.
-
-ONNX 1.20 removed several deprecated helper functions (float32_to_bfloat16,
-float32_to_float8e4m3, etc.). Some downstream packages like onnx_graphsurgeon 0.5.x
-still reference these at import time. This module restores them using ml_dtypes so
-that onnx_graphsurgeon can be imported without errors.
-
-This module must be imported BEFORE onnx_graphsurgeon.
-"""
-
-
-def patch_onnx_helper_removed_apis():
-    """Restore removed ONNX helper functions for backward compatibility."""
-    try:
-        import onnx.helper
-
-        if not hasattr(onnx.helper, "float32_to_bfloat16"):
-            import ml_dtypes
-            import numpy as np
-
-            def _float32_to_bfloat16(value, truncate=False):
-                return int.from_bytes(
-                    np.float32(value).astype(ml_dtypes.bfloat16).tobytes(), "little"
-                )
-
-            onnx.helper.float32_to_bfloat16 = _float32_to_bfloat16
-
-        if not hasattr(onnx.helper, "float32_to_float8e4m3"):
-            import ml_dtypes
-            import numpy as np
-
-            def _float32_to_float8e4m3(value, scale=1.0, fn=True, uz=False, saturate=True):
-                if fn and not uz:
-                    dtype = ml_dtypes.float8_e4m3fn
-                elif fn and uz:
-                    dtype = ml_dtypes.float8_e4m3fnuz
-                else:
-                    dtype = ml_dtypes.float8_e4m3fn
-                return int(np.float32(value).astype(dtype).view(np.uint8))
-
-            onnx.helper.float32_to_float8e4m3 = _float32_to_float8e4m3
-
-    except ImportError:
-        pass
-
-
-patch_onnx_helper_removed_apis()
diff --git a/pyproject.toml b/pyproject.toml
index fbe0575039c..61708763081 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -57,7 +57,7 @@ onnx = [
     "cupy-cuda12x; platform_machine != 'aarch64' and platform_system != 'Darwin'",
     "lief",
     "ml_dtypes",
-    "onnx-graphsurgeon",
+    "onnx-graphsurgeon>=0.6.1",
     "onnx~=1.21.0",
     "onnxconverter-common~=1.16.0",
     # ORT for Windows
diff --git a/tests/unit/onnx/conftest.py b/tests/unit/onnx/conftest.py
index f9e87a4706f..a08b2c2049a 100644
--- a/tests/unit/onnx/conftest.py
+++ b/tests/unit/onnx/conftest.py
@@ -12,6 +12,3 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-# Apply ONNX compatibility shim before test modules import onnx_graphsurgeon.
-from modelopt.onnx._onnx_compat import patch_onnx_helper_removed_apis  # noqa: F401
diff --git a/tests/unit/onnx/test_onnx_compat.py b/tests/unit/onnx/test_onnx_compat.py
deleted file mode 100644
index e16ada4e760..00000000000
--- a/tests/unit/onnx/test_onnx_compat.py
+++ /dev/null
@@ -1,131 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Tests for the ONNX compatibility shim (_onnx_compat)."""
-
-import importlib
-from unittest import mock
-
-import ml_dtypes
-import numpy as np
-import onnx.helper
-import pytest
-
-PATCHED_ATTRS = ["float32_to_bfloat16", "float32_to_float8e4m3"]
-
-
-def _reload_compat():
-    from modelopt.onnx import _onnx_compat
-
-    importlib.reload(_onnx_compat)
-
-
-@pytest.fixture
-def _clean_helper():
-    """Save, remove, and restore patched attrs around each test."""
-    saved = {attr: getattr(onnx.helper, attr, None) for attr in PATCHED_ATTRS}
-    for attr in PATCHED_ATTRS:
-        if hasattr(onnx.helper, attr):
-            delattr(onnx.helper, attr)
-    yield
-    for attr in PATCHED_ATTRS:
-        if saved[attr] is not None:
-            setattr(onnx.helper, attr, saved[attr])
-        elif hasattr(onnx.helper, attr):
-            delattr(onnx.helper, attr)
-
-
-@pytest.mark.usefixtures("_clean_helper")
-class TestPatchOnnxHelperRemovedApis:
-    @pytest.mark.parametrize("attr", PATCHED_ATTRS)
-    def test_patched_when_missing(self, attr):
-        _reload_compat()
-        assert callable(getattr(onnx.helper, attr))
-
-    @pytest.mark.parametrize("attr", PATCHED_ATTRS)
-    def test_existing_not_overwritten(self, attr):
-        sentinel = object()
-        setattr(onnx.helper, attr, sentinel)
-        _reload_compat()
-        assert getattr(onnx.helper, attr) is sentinel
-
-    def test_no_error_when_ml_dtypes_unavailable(self):
-        real_import = __import__
-
-        def _block_ml_dtypes(name, *args, **kwargs):
-            if name == "ml_dtypes":
-                raise ImportError("mocked")
-            return real_import(name, *args, **kwargs)
-
-        with mock.patch("builtins.__import__", side_effect=_block_ml_dtypes):
-            _reload_compat()
-
-
-@pytest.mark.parametrize(
-    ("value", "kwargs", "expected"),
-    [
-        (0.0, {}, 0),
-        (
-            1.0,
-            {},
-            int.from_bytes(np.float32(1.0).astype(ml_dtypes.bfloat16).tobytes(), "little"),
-        ),
-        (
-            1.0,
-            {"truncate": True},
-            int.from_bytes(np.float32(1.0).astype(ml_dtypes.bfloat16).tobytes(), "little"),
-        ),
-    ],
-)
-def test_bfloat16_conversion(value, kwargs, expected):
-    from modelopt.onnx import _onnx_compat  # noqa: F401
-
-    if not hasattr(onnx.helper, "float32_to_bfloat16"):
-        pytest.skip("ml_dtypes missing")
-    assert onnx.helper.float32_to_bfloat16(value, **kwargs) == expected
-
-
-@pytest.mark.parametrize(
-    ("value", "kwargs", "expected"),
-    [
-        (0.0, {}, 0),
-        (
-            1.0,
-            {},
-            int(np.float32(1.0).astype(ml_dtypes.float8_e4m3fn).view(np.uint8)),
-        ),
-        (
-            0.5,
-            {"fn": True, "uz": True},
-            int(np.float32(0.5).astype(ml_dtypes.float8_e4m3fnuz).view(np.uint8)),
-        ),
-        (
-            1.0,
-            {"scale": 2.0, "saturate": False},
-            int(np.float32(1.0).astype(ml_dtypes.float8_e4m3fn).view(np.uint8)),
-        ),
-        (
-            1.0,
-            {"fn": False},
-            int(np.float32(1.0).astype(ml_dtypes.float8_e4m3fn).view(np.uint8)),
-        ),
-    ],
-)
-def test_float8e4m3_conversion(value, kwargs, expected):
-    from modelopt.onnx import _onnx_compat  # noqa: F401
-
-    if not hasattr(onnx.helper, "float32_to_float8e4m3"):
-        pytest.skip("ml_dtypes missing")
-    assert onnx.helper.float32_to_float8e4m3(value, **kwargs) == expected

From 9eb8409e275cfbf7aa387351c53afcf316f38dc3 Mon Sep 17 00:00:00 2001
From: ajrasane <131806219+ajrasane@users.noreply.github.com>
Date: Thu, 9 Apr 2026 19:25:34 +0000
Subject: [PATCH 5/5] remove conftest.py

Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
---
 tests/unit/onnx/conftest.py | 14 --------------
 1 file changed, 14 deletions(-)
 delete mode 100644 tests/unit/onnx/conftest.py

diff --git a/tests/unit/onnx/conftest.py b/tests/unit/onnx/conftest.py
deleted file mode 100644
index a08b2c2049a..00000000000
--- a/tests/unit/onnx/conftest.py
+++ /dev/null
@@ -1,14 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.