From a22b2d4bc42380841cf0ff49de41b586a0a96cbc Mon Sep 17 00:00:00 2001 From: ConvolutedDog Date: Mon, 11 May 2026 18:36:17 +0800 Subject: [PATCH 1/4] [Fix][Relax]: ONNX Clip NaN bounds and preserve input NaN (ORT parity) --- .../tvm/relax/frontend/onnx/onnx_frontend.py | 36 +++++++++- tests/python/relax/test_frontend_onnx.py | 68 ++++++++++++++++++- 2 files changed, 98 insertions(+), 6 deletions(-) diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index 268d91b7500a..37ea8c6a5736 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -52,12 +52,26 @@ from tvm import TVMError, relax, tirx, topi from tvm.ir import IRModule from tvm.ir.supply import NameSupply +from tvm.runtime import DataType, DataTypeCode from tvm.tirx.generic import cast from tvm.topi.utils import get_const_tuple from ..common import autopad +def _relax_dtype_is_floating_point(dtype: str) -> bool: + """Whether a Relax dtype string is a floating point type.""" + try: + code = DataType(dtype).type_code + except (ValueError, TypeError, TVMError): + return False + return ( + code == DataTypeCode.FLOAT + or code == DataTypeCode.BFLOAT + or (code >= DataTypeCode.Float8E3M4 and code <= DataTypeCode.Float4E2M1FN) + ) + + def get_type(elem_type: str | int) -> str: """Converts onnx integer datatype to numpy datatype""" # If a string was passed instead of a tensor type, it does not need @@ -311,6 +325,7 @@ def get_converter(cls, opset): return getattr(cls, f"_impl_v{version}") raise NotImplementedError(f"opset version {version} of {cls.__name__} not implemented") + class QuantizeLinear(OnnxOpConverter): @classmethod def _impl_v10(cls, bb, inputs, attr, params): @@ -379,6 +394,7 @@ def _impl_v11(cls, bb, inputs, attr, params): y = relax.op.quantize(x, y_scale, y_zero_point, axis=0, out_dtype="uint8") return relax.Tuple([y, y_scale, y_zero_point]) + class MatMul(OnnxOpConverter): """Converts an onnx MatMul node into an equivalent Relax expression.""" @@ -1309,6 +1325,15 @@ def _impl_v16(cls, bb, inputs, attr, params): class Clip(OnnxOpConverter): """Converts an onnx Clip node into an equivalent Relax expression.""" + @staticmethod + def _sanitize_nan_clip_bound(bb, bound: relax.Expr, *, for_min: bool) -> relax.Expr: + """ONNX/ORT treat NaN clip bounds as unbounded; plain max/min with NaN poisons output.""" + dtype = bound.struct_info.dtype + if not _relax_dtype_is_floating_point(dtype): + return bound + repl = -_np.inf if for_min else _np.inf + return bb.emit(relax.op.where(relax.op.isnan(bound), relax.const(repl, dtype), bound)) + @classmethod def _impl_v1(cls, bb, inputs, attr, params): min = float(attr.get("min", -_np.inf)) @@ -1325,11 +1350,16 @@ def _impl_v11(cls, bb, inputs, attr, params): @classmethod def _impl_v13(cls, bb, inputs, attr, params): - results = inputs[0] + x: Any = inputs[0] + results = x if inputs[1] is not None: - results = bb.emit_te(topi.maximum, results, inputs[1]) + lo = cls._sanitize_nan_clip_bound(bb, inputs[1], for_min=True) + results = bb.emit_te(topi.maximum, results, lo) if inputs[2] is not None: - results = bb.emit_te(topi.minimum, results, inputs[2]) + hi = cls._sanitize_nan_clip_bound(bb, inputs[2], for_min=False) + results = bb.emit_te(topi.minimum, results, hi) + if _relax_dtype_is_floating_point(x.struct_info.dtype): + results = bb.emit(relax.op.where(relax.op.isnan(x), x, results)) return results diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index 5a8d84b0900c..a12479e2d3b6 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -90,6 +90,7 @@ def check_correctness( rtol: float = 1e-7, atol: float = 1e-5, check_dtypes: bool = False, + equal_nan: bool = False, ) -> None: """Run an onnx model in both onnxruntime and TVM through our importer confirm that the results match. Otherwise, an exception will be raised. @@ -109,6 +110,8 @@ def check_correctness( arithmetic variance than others. check_dtypes: bool Check if data types are the same. + equal_nan: bool + If True, treat NaN as equal when comparing floating outputs (np.testing.assert_allclose). """ # Configure model format. if ir_version is not None: @@ -177,16 +180,31 @@ def _check_output(tvm_out, ort_out): elif isinstance(tvm_out, tvm.runtime.Tensor) and isinstance(ort_out, np.ndarray): if check_dtypes: assert tvm_out.numpy().dtype == ort_out.dtype - tvm.testing.assert_allclose(tvm_out.numpy(), ort_out, rtol=rtol, atol=atol) + if equal_nan: + np.testing.assert_allclose( + tvm_out.numpy(), ort_out, rtol=rtol, atol=atol, equal_nan=True + ) + else: + tvm.testing.assert_allclose(tvm_out.numpy(), ort_out, rtol=rtol, atol=atol) elif isinstance(tvm_out, tvm_ffi.Shape) and isinstance(ort_out, np.ndarray): shape_out = tvm.runtime.tensor([int(i) for i in tvm_out]) if check_dtypes: assert _get_numpy_subdtype(shape_out.numpy()) == _get_numpy_subdtype(ort_out) - tvm.testing.assert_allclose(shape_out.numpy(), ort_out, rtol=rtol, atol=atol) + if equal_nan: + np.testing.assert_allclose( + shape_out.numpy(), ort_out, rtol=rtol, atol=atol, equal_nan=True + ) + else: + tvm.testing.assert_allclose(shape_out.numpy(), ort_out, rtol=rtol, atol=atol) elif isinstance(tvm_out, int | float | bool) and isinstance(ort_out, np.ndarray): if check_dtypes: assert _get_numpy_subdtype(np.array(tvm_out)) == _get_numpy_subdtype(ort_out) - tvm.testing.assert_allclose(np.array(tvm_out), ort_out, rtol=rtol, atol=atol) + if equal_nan: + np.testing.assert_allclose( + np.array(tvm_out), ort_out, rtol=rtol, atol=atol, equal_nan=True + ) + else: + tvm.testing.assert_allclose(np.array(tvm_out), ort_out, rtol=rtol, atol=atol) else: raise ValueError(f"Unsupported types: {type(tvm_out)}, {type(ort_out)}") @@ -1435,6 +1453,50 @@ def test_clip_v6(max, min): check_correctness(model, opset=10) +@pytest.mark.parametrize( + "min,max", + [ + pytest.param( + np.array(0.0, dtype=np.float32), + np.array(np.nan, dtype=np.float32), + ), + pytest.param( + np.array(0.0, dtype=np.float32), + np.array(np.nan, dtype=np.float32), + ), + pytest.param( + np.array(np.nan, dtype=np.float32), + np.array(6.0, dtype=np.float32), + ), + pytest.param( + np.array(np.nan, dtype=np.float32), + np.array(np.nan, dtype=np.float32), + ), + ], +) +def test_clip_v13(min, max): + # Opset 13: tensor min/max. NaN bound => unbounded on that side (ORT); input NaN preserved. + clip_node = helper.make_node("Clip", ["input", "min", "max"], ["output"]) + graph = helper.make_graph( + [clip_node], + "clip_v13_nan_max", + inputs=[ + helper.make_tensor_value_info("input", TensorProto.FLOAT, [5]), + helper.make_tensor_value_info("min", TensorProto.FLOAT, []), + helper.make_tensor_value_info("max", TensorProto.FLOAT, []), + ], + outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, [5])], + ) + model = helper.make_model(graph, producer_name="clip_v13_nan_max") + input = np.array([0.5, -3.0, 4.5, 11.0, 7.0], dtype=np.float32) + check_correctness( + model, + inputs={"input": input, "min": min, "max": max}, + opset=13, + equal_nan=True, + ) + + def test_equal(): equal_node = helper.make_node("Equal", ["a", "b"], ["output"]) From c6c1e49defca1aa6dd348c57595e1bf8926aa301 Mon Sep 17 00:00:00 2001 From: ConvolutedDog Date: Mon, 11 May 2026 20:18:13 +0800 Subject: [PATCH 2/4] [Feat][testing]: add equal_nan to assert_allclose; simplify ONNX check_correctness - Extend tvm.testing.assert_allclose with equal_nan (forwarded to NumPy) and document parameters. - Refactor check_correctness _check_output to normalize ndarray once and call a single assert helper. - Parametrize test_clip_v13 over input vectors (finite vs NaN); adjust Clip v13 min/max fixture where needed. --- python/tvm/testing/utils.py | 31 ++++++++++++-- tests/python/relax/test_frontend_onnx.py | 52 ++++++++++++------------ 2 files changed, 54 insertions(+), 29 deletions(-) diff --git a/python/tvm/testing/utils.py b/python/tvm/testing/utils.py index fa741f1d5c82..08683e72c4cc 100644 --- a/python/tvm/testing/utils.py +++ b/python/tvm/testing/utils.py @@ -103,18 +103,41 @@ def test_something(): ) -def assert_allclose(actual, desired, rtol=1e-7, atol=1e-7, verbose=True): - """Version of np.testing.assert_allclose with `atol` and `rtol` fields set - in reasonable defaults. +def assert_allclose(actual, desired, rtol=1e-7, atol=1e-7, verbose=True, equal_nan=False): + """Thin wrapper around `numpy.testing.assert_allclose` with TVM default tolerances. Arguments `actual` and `desired` are not interchangeable, since the function compares the `abs(actual-desired)` with `atol+rtol*abs(desired)`. Since we often allow `desired` to be close to zero, we generally want non-zero `atol`. + + Parameters + ---------- + actual : numpy.ndarray or array_like + Coerced with `numpy.asanyarray` before comparison. + + desired : numpy.ndarray or array_like + Reference values. These arguments are not interchangeable: the check uses + `atol + rtol * abs(desired)`. + + rtol : float, optional + Relative tolerance. Default is ``1e-7``. + + atol : float, optional + Absolute tolerance. Default is ``1e-7``. A non-zero value is usually + appropriate when `desired` may be close to zero. + + verbose : bool, optional + Passed through to NumPy. Default is ``True``. + + equal_nan : bool, optional + If True, NaNs at the same index are treated as equal. Default is ``False``. """ actual = np.asanyarray(actual) desired = np.asanyarray(desired) np.testing.assert_allclose(actual.shape, desired.shape) - np.testing.assert_allclose(actual, desired, rtol=rtol, atol=atol, verbose=verbose) + np.testing.assert_allclose( + actual, desired, rtol=rtol, atol=atol, verbose=verbose, equal_nan=equal_nan + ) def check_numerical_grads( diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index a12479e2d3b6..a9bb1d244ac4 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -173,38 +173,34 @@ def _get_numpy_subdtype(narray): return "other" def _check_output(tvm_out, ort_out): + def _assert_allclose_nd(actual_np: np.ndarray, desired_np: np.ndarray) -> None: + tvm.testing.assert_allclose( + actual_np, + desired_np, + rtol=rtol, + atol=atol, + equal_nan=equal_nan, + ) + if isinstance(tvm_out, tuple) and isinstance(ort_out, tvm_ffi.Shape | list): assert len(tvm_out) == len(ort_out), "Unequal number of outputs" for tvm_out_i, ort_out_i in zip(tvm_out, ort_out): _check_output(tvm_out_i, ort_out_i) elif isinstance(tvm_out, tvm.runtime.Tensor) and isinstance(ort_out, np.ndarray): + actual_np = tvm_out.numpy() if check_dtypes: - assert tvm_out.numpy().dtype == ort_out.dtype - if equal_nan: - np.testing.assert_allclose( - tvm_out.numpy(), ort_out, rtol=rtol, atol=atol, equal_nan=True - ) - else: - tvm.testing.assert_allclose(tvm_out.numpy(), ort_out, rtol=rtol, atol=atol) + assert actual_np.dtype == ort_out.dtype + _assert_allclose_nd(actual_np, ort_out) elif isinstance(tvm_out, tvm_ffi.Shape) and isinstance(ort_out, np.ndarray): - shape_out = tvm.runtime.tensor([int(i) for i in tvm_out]) + actual_np = tvm.runtime.tensor([int(i) for i in tvm_out]).numpy() if check_dtypes: - assert _get_numpy_subdtype(shape_out.numpy()) == _get_numpy_subdtype(ort_out) - if equal_nan: - np.testing.assert_allclose( - shape_out.numpy(), ort_out, rtol=rtol, atol=atol, equal_nan=True - ) - else: - tvm.testing.assert_allclose(shape_out.numpy(), ort_out, rtol=rtol, atol=atol) + assert _get_numpy_subdtype(actual_np) == _get_numpy_subdtype(ort_out) + _assert_allclose_nd(actual_np, ort_out) elif isinstance(tvm_out, int | float | bool) and isinstance(ort_out, np.ndarray): + actual_np = np.array(tvm_out) if check_dtypes: - assert _get_numpy_subdtype(np.array(tvm_out)) == _get_numpy_subdtype(ort_out) - if equal_nan: - np.testing.assert_allclose( - np.array(tvm_out), ort_out, rtol=rtol, atol=atol, equal_nan=True - ) - else: - tvm.testing.assert_allclose(np.array(tvm_out), ort_out, rtol=rtol, atol=atol) + assert _get_numpy_subdtype(actual_np) == _get_numpy_subdtype(ort_out) + _assert_allclose_nd(actual_np, ort_out) else: raise ValueError(f"Unsupported types: {type(tvm_out)}, {type(ort_out)}") @@ -1458,7 +1454,7 @@ def test_clip_v6(max, min): [ pytest.param( np.array(0.0, dtype=np.float32), - np.array(np.nan, dtype=np.float32), + np.array(6.0, dtype=np.float32), ), pytest.param( np.array(0.0, dtype=np.float32), @@ -1474,7 +1470,14 @@ def test_clip_v6(max, min): ), ], ) -def test_clip_v13(min, max): +@pytest.mark.parametrize( + "input", + [ + np.array([0.5, -3.0, 4.5, 11.0, 7.0], dtype=np.float32), + np.array([0.5, -3.0, 4.5, 11.0, np.nan], dtype=np.float32), + ], +) +def test_clip_v13(input, min, max): # Opset 13: tensor min/max. NaN bound => unbounded on that side (ORT); input NaN preserved. clip_node = helper.make_node("Clip", ["input", "min", "max"], ["output"]) graph = helper.make_graph( @@ -1488,7 +1491,6 @@ def test_clip_v13(min, max): outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, [5])], ) model = helper.make_model(graph, producer_name="clip_v13_nan_max") - input = np.array([0.5, -3.0, 4.5, 11.0, 7.0], dtype=np.float32) check_correctness( model, inputs={"input": input, "min": min, "max": max}, From e61617c4b997fa93dd504bbe9c0acd75f73b1ab2 Mon Sep 17 00:00:00 2001 From: ConvolutedDog Date: Tue, 12 May 2026 00:11:39 +0800 Subject: [PATCH 3/4] revert --- python/tvm/testing/utils.py | 31 +++---------------- tests/python/relax/test_frontend_onnx.py | 25 +++++---------- .../test_meta_schedule_search_strategy.py | 2 +- 3 files changed, 12 insertions(+), 46 deletions(-) diff --git a/python/tvm/testing/utils.py b/python/tvm/testing/utils.py index 08683e72c4cc..fa741f1d5c82 100644 --- a/python/tvm/testing/utils.py +++ b/python/tvm/testing/utils.py @@ -103,41 +103,18 @@ def test_something(): ) -def assert_allclose(actual, desired, rtol=1e-7, atol=1e-7, verbose=True, equal_nan=False): - """Thin wrapper around `numpy.testing.assert_allclose` with TVM default tolerances. +def assert_allclose(actual, desired, rtol=1e-7, atol=1e-7, verbose=True): + """Version of np.testing.assert_allclose with `atol` and `rtol` fields set + in reasonable defaults. Arguments `actual` and `desired` are not interchangeable, since the function compares the `abs(actual-desired)` with `atol+rtol*abs(desired)`. Since we often allow `desired` to be close to zero, we generally want non-zero `atol`. - - Parameters - ---------- - actual : numpy.ndarray or array_like - Coerced with `numpy.asanyarray` before comparison. - - desired : numpy.ndarray or array_like - Reference values. These arguments are not interchangeable: the check uses - `atol + rtol * abs(desired)`. - - rtol : float, optional - Relative tolerance. Default is ``1e-7``. - - atol : float, optional - Absolute tolerance. Default is ``1e-7``. A non-zero value is usually - appropriate when `desired` may be close to zero. - - verbose : bool, optional - Passed through to NumPy. Default is ``True``. - - equal_nan : bool, optional - If True, NaNs at the same index are treated as equal. Default is ``False``. """ actual = np.asanyarray(actual) desired = np.asanyarray(desired) np.testing.assert_allclose(actual.shape, desired.shape) - np.testing.assert_allclose( - actual, desired, rtol=rtol, atol=atol, verbose=verbose, equal_nan=equal_nan - ) + np.testing.assert_allclose(actual, desired, rtol=rtol, atol=atol, verbose=verbose) def check_numerical_grads( diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index a9bb1d244ac4..8aae062f02f1 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -173,34 +173,23 @@ def _get_numpy_subdtype(narray): return "other" def _check_output(tvm_out, ort_out): - def _assert_allclose_nd(actual_np: np.ndarray, desired_np: np.ndarray) -> None: - tvm.testing.assert_allclose( - actual_np, - desired_np, - rtol=rtol, - atol=atol, - equal_nan=equal_nan, - ) - if isinstance(tvm_out, tuple) and isinstance(ort_out, tvm_ffi.Shape | list): assert len(tvm_out) == len(ort_out), "Unequal number of outputs" for tvm_out_i, ort_out_i in zip(tvm_out, ort_out): _check_output(tvm_out_i, ort_out_i) elif isinstance(tvm_out, tvm.runtime.Tensor) and isinstance(ort_out, np.ndarray): - actual_np = tvm_out.numpy() if check_dtypes: - assert actual_np.dtype == ort_out.dtype - _assert_allclose_nd(actual_np, ort_out) + assert tvm_out.numpy().dtype == ort_out.dtype + tvm.testing.assert_allclose(tvm_out.numpy(), ort_out, rtol=rtol, atol=atol) elif isinstance(tvm_out, tvm_ffi.Shape) and isinstance(ort_out, np.ndarray): - actual_np = tvm.runtime.tensor([int(i) for i in tvm_out]).numpy() + shape_out = tvm.runtime.tensor([int(i) for i in tvm_out]) if check_dtypes: - assert _get_numpy_subdtype(actual_np) == _get_numpy_subdtype(ort_out) - _assert_allclose_nd(actual_np, ort_out) + assert _get_numpy_subdtype(shape_out.numpy()) == _get_numpy_subdtype(ort_out) + tvm.testing.assert_allclose(shape_out.numpy(), ort_out, rtol=rtol, atol=atol) elif isinstance(tvm_out, int | float | bool) and isinstance(ort_out, np.ndarray): - actual_np = np.array(tvm_out) if check_dtypes: - assert _get_numpy_subdtype(actual_np) == _get_numpy_subdtype(ort_out) - _assert_allclose_nd(actual_np, ort_out) + assert _get_numpy_subdtype(np.array(tvm_out)) == _get_numpy_subdtype(ort_out) + tvm.testing.assert_allclose(np.array(tvm_out), ort_out, rtol=rtol, atol=atol) else: raise ValueError(f"Unsupported types: {type(tvm_out)}, {type(ort_out)}") diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_search_strategy.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_search_strategy.py index 370eff27c77b..f9cec06aea9d 100644 --- a/tests/python/s_tir/meta_schedule/test_meta_schedule_search_strategy.py +++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_search_strategy.py @@ -324,7 +324,7 @@ def __str__(self) -> str: assert candidates is None -def test_meta_schedule_evolutionary_search_skip_invalid_measured_trace() # pylint: disable = invalid-name +def test_meta_schedule_evolutionary_search_skip_invalid_measured_trace(): # pylint: disable = invalid-name # Construct an incompatible measured trace: it references block name "other", # which doesn't exist in Matmul. Replaying this trace should fail and be skipped. wrong_sch = Schedule(OtherBlock) From f08f3c7d555c47e27c00b72ddb65fb0d916964d2 Mon Sep 17 00:00:00 2001 From: ConvolutedDog Date: Tue, 12 May 2026 15:13:49 +0800 Subject: [PATCH 4/4] remove equal_nan parameter --- tests/python/relax/test_frontend_onnx.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index 8aae062f02f1..2c9cf1e75c2b 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -90,7 +90,6 @@ def check_correctness( rtol: float = 1e-7, atol: float = 1e-5, check_dtypes: bool = False, - equal_nan: bool = False, ) -> None: """Run an onnx model in both onnxruntime and TVM through our importer confirm that the results match. Otherwise, an exception will be raised. @@ -110,8 +109,6 @@ def check_correctness( arithmetic variance than others. check_dtypes: bool Check if data types are the same. - equal_nan: bool - If True, treat NaN as equal when comparing floating outputs (np.testing.assert_allclose). """ # Configure model format. if ir_version is not None: @@ -1484,7 +1481,6 @@ def test_clip_v13(input, min, max): model, inputs={"input": input, "min": min, "max": max}, opset=13, - equal_nan=True, )