From 79afc041f22060a1a9c70f413d0fafecb5fbc86e Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Sun, 28 Jun 2026 20:04:32 +0800 Subject: [PATCH 1/5] [Relax][ONNX] Add CastLike support and unblock Relu backend tests - Add a CastLike converter to the Relax ONNX frontend (opset 15+). - CastLike casts the first input to the dtype of the second input, which is required by the expanded form of Relu in opset 18. - Add relu to the backend test allowlist. - Add a focused CastLike regression test. Relates to issue #19505. --- python/tvm/relax/frontend/onnx/onnx_frontend.py | 12 ++++++++++++ tests/python/relax/test_frontend_onnx.py | 15 +++++++++++++++ tests/python/relax/test_frontend_onnx_backend.py | 12 +++++++++--- 3 files changed, 36 insertions(+), 3 deletions(-) diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index 87675a3f9c12..d1376aa2df4b 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -1170,6 +1170,17 @@ def _impl_v13(cls, bb, inputs, attr, params): return relax.op.astype(inputs[0], to_type) +class CastLike(OnnxOpConverter): + """Convert an onnx CastLike node into an equivalent Relax expression.""" + + @classmethod + def _impl_v15(cls, bb, inputs, attr, params): + data = inputs[0] + target = inputs[1] + target_dtype = target.ty.dtype.dtype + return relax.op.astype(data, target_dtype) + + class Gather(OnnxOpConverter): """Convert an onnx Gather node into an equivalent Relax expression.""" @@ -5241,6 +5252,7 @@ def _get_convert_map(): "Max": Max, "Mean": Mean, "Cast": Cast, + "CastLike": CastLike, "Gemm": Gemm, "MatMul": MatMul, "MatMulInteger": MatMulInteger, diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index 972bc4830729..ea181bd87597 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -1358,6 +1358,21 @@ def test_cast_nan_inf_to_int8(): np.testing.assert_array_equal(out_np, expected) +def test_castlike(): + castlike_node = helper.make_node("CastLike", ["a", "b"], ["c"]) + graph = helper.make_graph( + [castlike_node], + "castlike_test", + inputs=[ + helper.make_tensor_value_info("a", TensorProto.FLOAT, [1, 32]), + helper.make_tensor_value_info("b", TensorProto.INT32, [1]), + ], + outputs=[helper.make_tensor_value_info("c", TensorProto.INT32, [1, 32])], + ) + model = helper.make_model(graph, producer_name="castlike_test") + check_correctness(model, opset=15, check_dtypes=True) + + def test_gather(): def _verify_gather(data_shape, indices, out_shape, expected, axis=0): gather_node = helper.make_node("Gather", ["data", "indices"], ["y"], axis=axis) diff --git a/tests/python/relax/test_frontend_onnx_backend.py b/tests/python/relax/test_frontend_onnx_backend.py index 9f9ec8779cdd..c47eba898e7e 100644 --- a/tests/python/relax/test_frontend_onnx_backend.py +++ b/tests/python/relax/test_frontend_onnx_backend.py @@ -81,9 +81,9 @@ def run(self, inputs, **kwargs): self._vm.invoke_stateful("main") output = self._vm.get_outputs("main") - if isinstance(output, (tvm.runtime.Tensor, np.ndarray)): # noqa: UP038 + if isinstance(output, (tvm.runtime.Tensor, np.ndarray)): return (output.numpy() if hasattr(output, "numpy") else output,) - if isinstance(output, (tuple, list)): # noqa: UP038 + if isinstance(output, (tuple, list)): return tuple(o.numpy() if hasattr(o, "numpy") else np.array(o) for o in output) return (np.array(output),) @@ -172,6 +172,7 @@ def supports_device(cls, device: str) -> bool: "less", "less_equal", "lrn", + "logsoftmax", "matmul", "matmulinteger", "mean", @@ -183,6 +184,7 @@ def supports_device(cls, device: str) -> bool: "not", "or", "reciprocal", + "relu", "round", "scatternd", "sigmoid", @@ -191,6 +193,7 @@ def supports_device(cls, device: str) -> bool: "sinh", "size", "slice", + "softmax", "spacetodepth", "sqrt", "squeeze", @@ -209,4 +212,7 @@ def supports_device(cls, device: str) -> bool: for _op in _INCLUDE_OPS: backend_test.include(rf"^test_{_op}(?:_.*)?(?:_cpu|_cuda)$") -globals().update(backend_test.test_cases) +# Only node-level backend tests are in scope for importer conformance. +globals().update( + {k: v for k, v in backend_test.test_cases.items() if k == "OnnxBackendNodeModelTest"} +) From 2b1cba2e323bc2df08391eb9efd47198e70299cd Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Sun, 28 Jun 2026 20:18:23 +0800 Subject: [PATCH 2/5] [Relax][ONNX] Support dynamic k in Trilu and expand backend allowlist - Remove the constant-k restriction from the Trilu converter. - For constant k, keep using the optimized relax.op.tril/triu. - For dynamic k, build the lower/upper-triangular mask explicitly with arange and broadcast it to the input shape. - Add tril and triu to the backend test allowlist. - Add a regression test for Trilu with dynamic k. Relates to issue #19505. --- .../tvm/relax/frontend/onnx/onnx_frontend.py | 30 ++++-- tests/python/relax/test_frontend_onnx.py | 91 ++++++++++++++++++- .../relax/test_frontend_onnx_backend.py | 2 + 3 files changed, 111 insertions(+), 12 deletions(-) diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index d1376aa2df4b..187f6a10cd8a 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -1517,19 +1517,29 @@ def _impl_v14(cls, bb, inputs, attr, params): x = inputs[0] k = inputs[1] if len(inputs) > 1 else 0 - if len(inputs) > 1: - k = get_constant(inputs[1], params) - if isinstance(k, relax.Constant): - k = int(k.data.numpy().item()) - else: - raise ValueError("Currently only support constant k for Trilu op.") - else: - k = 0 + if isinstance(k, relax.Constant): + k = int(k.data.numpy().item()) + if isinstance(k, int): + if upper: + return relax.op.triu(x, k) + return relax.op.tril(x, k) + # Dynamic k: build the mask explicitly so it works with any scalar k. + shape = x.ty.shape + m, n = shape[-2], shape[-1] + row_idx = relax.op.reshape(relax.op.arange(0, m, dtype="int64"), (m, 1)) + col_idx = relax.op.reshape(relax.op.arange(0, n, dtype="int64"), (1, n)) + diff = relax.op.subtract( + relax.op.broadcast_to(col_idx, (m, n)), + relax.op.broadcast_to(row_idx, (m, n)), + ) + k_int64 = relax.op.astype(k, "int64") if upper: - return relax.op.triu(x, k) + mask = relax.op.greater_equal(diff, k_int64) else: - return relax.op.tril(x, k) + mask = relax.op.less_equal(diff, k_int64) + mask = relax.op.broadcast_to(mask, shape) + return relax.op.where(mask, x, relax.const(0, x.ty.dtype)) class Relu(OnnxOpConverter): diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index ea181bd87597..89c21a0f1c2e 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -1358,7 +1358,7 @@ def test_cast_nan_inf_to_int8(): np.testing.assert_array_equal(out_np, expected) -def test_castlike(): +def test_castlike_ir(): castlike_node = helper.make_node("CastLike", ["a", "b"], ["c"]) graph = helper.make_graph( [castlike_node], @@ -1370,7 +1370,22 @@ def test_castlike(): outputs=[helper.make_tensor_value_info("c", TensorProto.INT32, [1, 32])], ) model = helper.make_model(graph, producer_name="castlike_test") - check_correctness(model, opset=15, check_dtypes=True) + tvm_model = from_onnx(model, opset=15, keep_params_in_input=True) + + @I.ir_module + class Expected: + @R.function + def main( + a: R.Tensor((1, 32), dtype="float32"), + b: R.Tensor((1,), dtype="int32"), + ) -> R.Tensor((1, 32), dtype="int32"): + R.func_attr({"num_input": 2}) + with R.dataflow(): + gv: R.Tensor((1, 32), dtype="int32") = R.astype(a, "int32") + R.output(gv) + return gv + + tvm.ir.assert_structural_equal(tvm_model, Expected) def test_gather(): @@ -3103,6 +3118,78 @@ def test_trilu_with_const_k(k_value: int): check_correctness(model) +@pytest.mark.parametrize("upper", [True, False]) +def test_trilu_dynamic_k_ir(upper: bool): + graph = helper.make_graph( + [ + helper.make_node("Trilu", inputs=["x", "k"], outputs=["y"], upper=upper), + ], + "trilu_dynamic_k_graph", + inputs=[ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3]), + helper.make_tensor_value_info("k", TensorProto.INT64, []), + ], + outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])], + ) + model = helper.make_model(graph, producer_name="trilu_dynamic_k_graph") + tvm_model = from_onnx(model, opset=14, keep_params_in_input=True) + + if upper: + + @I.ir_module + class ExpectedTriu: + @R.function + def main( + x: R.Tensor((2, 3), dtype="float32"), + k: R.Tensor((), dtype="int64"), + ) -> R.Tensor((2, 3), dtype="float32"): + R.func_attr({"num_input": 2}) + with R.dataflow(): + lv: R.Tensor((3,), dtype="int64") = R.arange(0, 3, 1, dtype="int64") + lv1: R.Tensor((1, 3), dtype="int64") = R.reshape(lv, R.shape([1, 3])) + lv2: R.Tensor((2, 3), dtype="int64") = R.broadcast_to(lv1, R.shape([2, 3])) + lv3: R.Tensor((2,), dtype="int64") = R.arange(0, 2, 1, dtype="int64") + lv4: R.Tensor((2, 1), dtype="int64") = R.reshape(lv3, R.shape([2, 1])) + lv5: R.Tensor((2, 3), dtype="int64") = R.broadcast_to(lv4, R.shape([2, 3])) + lv6: R.Tensor((2, 3), dtype="int64") = R.subtract(lv2, lv5) + lv7: R.Tensor((), dtype="int64") = R.astype(k, dtype="int64") + lv8: R.Tensor((2, 3), dtype="bool") = R.greater_equal(lv6, lv7) + lv9: R.Tensor((2, 3), dtype="bool") = R.broadcast_to(lv8, R.shape([2, 3])) + gv: R.Tensor((2, 3), dtype="float32") = R.where(lv9, x, R.const(0.0, "float32")) + R.output(gv) + return gv + + expected = ExpectedTriu + else: + + @I.ir_module + class ExpectedTril: + @R.function + def main( + x: R.Tensor((2, 3), dtype="float32"), + k: R.Tensor((), dtype="int64"), + ) -> R.Tensor((2, 3), dtype="float32"): + R.func_attr({"num_input": 2}) + with R.dataflow(): + lv: R.Tensor((3,), dtype="int64") = R.arange(0, 3, 1, dtype="int64") + lv1: R.Tensor((1, 3), dtype="int64") = R.reshape(lv, R.shape([1, 3])) + lv2: R.Tensor((2, 3), dtype="int64") = R.broadcast_to(lv1, R.shape([2, 3])) + lv3: R.Tensor((2,), dtype="int64") = R.arange(0, 2, 1, dtype="int64") + lv4: R.Tensor((2, 1), dtype="int64") = R.reshape(lv3, R.shape([2, 1])) + lv5: R.Tensor((2, 3), dtype="int64") = R.broadcast_to(lv4, R.shape([2, 3])) + lv6: R.Tensor((2, 3), dtype="int64") = R.subtract(lv2, lv5) + lv7: R.Tensor((), dtype="int64") = R.astype(k, dtype="int64") + lv8: R.Tensor((2, 3), dtype="bool") = R.less_equal(lv6, lv7) + lv9: R.Tensor((2, 3), dtype="bool") = R.broadcast_to(lv8, R.shape([2, 3])) + gv: R.Tensor((2, 3), dtype="float32") = R.where(lv9, x, R.const(0.0, "float32")) + R.output(gv) + return gv + + expected = ExpectedTril + + tvm.ir.assert_structural_equal(tvm_model, expected) + + def test_selu(): model = make_unary_model("Selu", [2, 3]) tvm_model = from_onnx(model, keep_params_in_input=True) diff --git a/tests/python/relax/test_frontend_onnx_backend.py b/tests/python/relax/test_frontend_onnx_backend.py index c47eba898e7e..2605fbaddcda 100644 --- a/tests/python/relax/test_frontend_onnx_backend.py +++ b/tests/python/relax/test_frontend_onnx_backend.py @@ -203,6 +203,8 @@ def supports_device(cls, device: str) -> bool: "tanh", "tile", "transpose", + "tril", + "triu", "unique", "unsqueeze", "where", From dab4f7486f822e5839f70282b32472aa2f8f633b Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Sun, 28 Jun 2026 21:22:55 +0800 Subject: [PATCH 3/5] [Relax][ONNX] Keep backend test registration open for model-level extensions --- .../relax/test_frontend_onnx_backend.py | 28 ++++++++++++------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/tests/python/relax/test_frontend_onnx_backend.py b/tests/python/relax/test_frontend_onnx_backend.py index 2605fbaddcda..ef61bd203484 100644 --- a/tests/python/relax/test_frontend_onnx_backend.py +++ b/tests/python/relax/test_frontend_onnx_backend.py @@ -19,13 +19,13 @@ ONNX Backend Tests =================== Systematically verify the Relax ONNX importer using the official ONNX -Backend Test Suite (node-level tests only). Each test loads a small -ONNX model with protobuf reference inputs/outputs and checks that the -Relax-imported model produces numerically correct results. +Backend Test Suite. Each test loads a small ONNX model with protobuf +reference inputs/outputs and checks that the Relax-imported model +produces numerically correct results. -Only ``onnx.backend.test.data.node`` tests are registered here; real, -simple, and PyTorch model tests are out of scope for importer-level -semantic verification. +Currently ``_INCLUDE_OPS`` selects node-level operator tests. Other +test classes (real/simple/PyTorch models) remain available in +``backend_test.test_cases`` and can be enabled explicitly in the future. """ @@ -214,7 +214,15 @@ def supports_device(cls, device: str) -> bool: for _op in _INCLUDE_OPS: backend_test.include(rf"^test_{_op}(?:_.*)?(?:_cpu|_cuda)$") -# Only node-level backend tests are in scope for importer conformance. -globals().update( - {k: v for k, v in backend_test.test_cases.items() if k == "OnnxBackendNodeModelTest"} -) +# A small number of model-level tests (e.g. from PyTorch converted models) +# have names that collide with the node-level include patterns above. The +# current adapter is focused on node-level protobuf test cases, so exclude +# those known collisions explicitly rather than limiting the test classes. +_EXCLUDE_PATTERNS = [ + r"^test_softmax_functional_dim3_cpu$", + r"^test_softmax_lastdim_cpu$", +] +for _pattern in _EXCLUDE_PATTERNS: + backend_test.exclude(_pattern) + +globals().update(backend_test.test_cases) From 2acccfb4c2d7f476c63deade2c31ee0288f36dc7 Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Sun, 28 Jun 2026 21:35:50 +0800 Subject: [PATCH 4/5] [Relax][ONNX] Fix UP038 lint warnings in backend test adapter --- tests/python/relax/test_frontend_onnx_backend.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/python/relax/test_frontend_onnx_backend.py b/tests/python/relax/test_frontend_onnx_backend.py index ef61bd203484..2a6c7bdbcdcf 100644 --- a/tests/python/relax/test_frontend_onnx_backend.py +++ b/tests/python/relax/test_frontend_onnx_backend.py @@ -81,9 +81,9 @@ def run(self, inputs, **kwargs): self._vm.invoke_stateful("main") output = self._vm.get_outputs("main") - if isinstance(output, (tvm.runtime.Tensor, np.ndarray)): + if isinstance(output, tvm.runtime.Tensor | np.ndarray): return (output.numpy() if hasattr(output, "numpy") else output,) - if isinstance(output, (tuple, list)): + if isinstance(output, tuple | list): return tuple(o.numpy() if hasattr(o, "numpy") else np.array(o) for o in output) return (np.array(output),) From a73ff50d858e8ecf4f83875a4e9673a54562434d Mon Sep 17 00:00:00 2001 From: Aharrypotter Date: Sun, 28 Jun 2026 22:53:40 +0800 Subject: [PATCH 5/5] [Relax][ONNX] Apply review feedback to CastLike and Trilu --- .../tvm/relax/frontend/onnx/onnx_frontend.py | 13 ++++---- tests/python/relax/test_frontend_onnx.py | 32 ++++++++----------- 2 files changed, 21 insertions(+), 24 deletions(-) diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index 187f6a10cd8a..855972b9489d 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -1177,7 +1177,11 @@ class CastLike(OnnxOpConverter): def _impl_v15(cls, bb, inputs, attr, params): data = inputs[0] target = inputs[1] - target_dtype = target.ty.dtype.dtype + target_dtype = getattr(getattr(getattr(target, "ty", None), "dtype", None), "dtype", None) + if target_dtype is None: + target_dtype = getattr(target.struct_info, "dtype", None) + if target_dtype is None: + raise ValueError(f"CastLike: unable to determine dtype from target {target}") return relax.op.astype(data, target_dtype) @@ -1529,17 +1533,14 @@ def _impl_v14(cls, bb, inputs, attr, params): m, n = shape[-2], shape[-1] row_idx = relax.op.reshape(relax.op.arange(0, m, dtype="int64"), (m, 1)) col_idx = relax.op.reshape(relax.op.arange(0, n, dtype="int64"), (1, n)) - diff = relax.op.subtract( - relax.op.broadcast_to(col_idx, (m, n)), - relax.op.broadcast_to(row_idx, (m, n)), - ) + diff = relax.op.subtract(col_idx, row_idx) k_int64 = relax.op.astype(k, "int64") if upper: mask = relax.op.greater_equal(diff, k_int64) else: mask = relax.op.less_equal(diff, k_int64) mask = relax.op.broadcast_to(mask, shape) - return relax.op.where(mask, x, relax.const(0, x.ty.dtype)) + return relax.op.where(mask, x, relax.const(0, x.ty.dtype.dtype)) class Relu(OnnxOpConverter): diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index 89c21a0f1c2e..0ad958c62d6c 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -3147,15 +3147,13 @@ def main( with R.dataflow(): lv: R.Tensor((3,), dtype="int64") = R.arange(0, 3, 1, dtype="int64") lv1: R.Tensor((1, 3), dtype="int64") = R.reshape(lv, R.shape([1, 3])) - lv2: R.Tensor((2, 3), dtype="int64") = R.broadcast_to(lv1, R.shape([2, 3])) - lv3: R.Tensor((2,), dtype="int64") = R.arange(0, 2, 1, dtype="int64") - lv4: R.Tensor((2, 1), dtype="int64") = R.reshape(lv3, R.shape([2, 1])) - lv5: R.Tensor((2, 3), dtype="int64") = R.broadcast_to(lv4, R.shape([2, 3])) - lv6: R.Tensor((2, 3), dtype="int64") = R.subtract(lv2, lv5) - lv7: R.Tensor((), dtype="int64") = R.astype(k, dtype="int64") - lv8: R.Tensor((2, 3), dtype="bool") = R.greater_equal(lv6, lv7) - lv9: R.Tensor((2, 3), dtype="bool") = R.broadcast_to(lv8, R.shape([2, 3])) - gv: R.Tensor((2, 3), dtype="float32") = R.where(lv9, x, R.const(0.0, "float32")) + lv2: R.Tensor((2,), dtype="int64") = R.arange(0, 2, 1, dtype="int64") + lv3: R.Tensor((2, 1), dtype="int64") = R.reshape(lv2, R.shape([2, 1])) + lv4: R.Tensor((2, 3), dtype="int64") = R.subtract(lv1, lv3) + lv5: R.Tensor((), dtype="int64") = R.astype(k, dtype="int64") + lv6: R.Tensor((2, 3), dtype="bool") = R.greater_equal(lv4, lv5) + lv7: R.Tensor((2, 3), dtype="bool") = R.broadcast_to(lv6, R.shape([2, 3])) + gv: R.Tensor((2, 3), dtype="float32") = R.where(lv7, x, R.const(0.0, "float32")) R.output(gv) return gv @@ -3173,15 +3171,13 @@ def main( with R.dataflow(): lv: R.Tensor((3,), dtype="int64") = R.arange(0, 3, 1, dtype="int64") lv1: R.Tensor((1, 3), dtype="int64") = R.reshape(lv, R.shape([1, 3])) - lv2: R.Tensor((2, 3), dtype="int64") = R.broadcast_to(lv1, R.shape([2, 3])) - lv3: R.Tensor((2,), dtype="int64") = R.arange(0, 2, 1, dtype="int64") - lv4: R.Tensor((2, 1), dtype="int64") = R.reshape(lv3, R.shape([2, 1])) - lv5: R.Tensor((2, 3), dtype="int64") = R.broadcast_to(lv4, R.shape([2, 3])) - lv6: R.Tensor((2, 3), dtype="int64") = R.subtract(lv2, lv5) - lv7: R.Tensor((), dtype="int64") = R.astype(k, dtype="int64") - lv8: R.Tensor((2, 3), dtype="bool") = R.less_equal(lv6, lv7) - lv9: R.Tensor((2, 3), dtype="bool") = R.broadcast_to(lv8, R.shape([2, 3])) - gv: R.Tensor((2, 3), dtype="float32") = R.where(lv9, x, R.const(0.0, "float32")) + lv2: R.Tensor((2,), dtype="int64") = R.arange(0, 2, 1, dtype="int64") + lv3: R.Tensor((2, 1), dtype="int64") = R.reshape(lv2, R.shape([2, 1])) + lv4: R.Tensor((2, 3), dtype="int64") = R.subtract(lv1, lv3) + lv5: R.Tensor((), dtype="int64") = R.astype(k, dtype="int64") + lv6: R.Tensor((2, 3), dtype="bool") = R.less_equal(lv4, lv5) + lv7: R.Tensor((2, 3), dtype="bool") = R.broadcast_to(lv6, R.shape([2, 3])) + gv: R.Tensor((2, 3), dtype="float32") = R.where(lv7, x, R.const(0.0, "float32")) R.output(gv) return gv