From 9b2f448dc94e76cdfabfdb8910895ab959a1209a Mon Sep 17 00:00:00 2001 From: viiccwen Date: Fri, 10 Jul 2026 16:13:10 +0000 Subject: [PATCH 1/2] [Fix][Relax][ONNX] Fix Range dtype for PrimExpr bounds Signed-off-by: viiccwen --- .../tvm/relax/frontend/onnx/onnx_frontend.py | 2 +- tests/python/relax/test_frontend_onnx.py | 30 +++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index 0ca7ef9e0c0a..95c39b5d7029 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -3626,7 +3626,7 @@ def _impl_v12(cls, bb, inputs, attr, params): start = get_constant(inputs[0], params) limit = get_constant(inputs[1], params) delta = get_constant(inputs[2], params) - out_dtype = start.ty.dtype + out_dtype = start.ty.dtype.dtype if isinstance(start, relax.Constant): start = start.data.numpy().tolist() diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index a21b3d4e7a8d..d1505e7360ff 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -7760,6 +7760,36 @@ def main( tvm.ir.assert_structural_equal(tvm_model, Expected) +def test_range_with_primexpr_limit(): + shape = helper.make_node("Shape", ["x"], ["x_shape"]) + axis = make_constant_node("axis", TensorProto.INT64, [], [0]) + limit = helper.make_node("Gather", ["x_shape", "axis"], ["limit"]) + start = make_constant_node("start", TensorProto.INT64, [], [0]) + delta = make_constant_node("delta", TensorProto.INT64, [], [1]) + range_node = helper.make_node("Range", ["start", "limit", "delta"], ["output"]) + graph = helper.make_graph( + [shape, axis, limit, start, delta, range_node], + "range_primexpr_limit_test", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [5])], + [helper.make_tensor_value_info("output", TensorProto.INT64, [5])], + ) + model = helper.make_model(graph, producer_name="range_primexpr_limit_test") + + tvm_model = from_onnx(model) + + @I.ir_module + class Expected: + @R.function + def main(x: R.Tensor((5,), dtype="float32")) -> R.Tensor((5,), dtype="int64"): + R.func_attr({"num_input": 1}) + with R.dataflow(): + gv: R.Tensor((5,), dtype="int64") = R.arange(0, 5, 1, dtype="int64") + R.output(gv) + return gv + + tvm.ir.assert_structural_equal(tvm_model, Expected) + + def test_batch_norm(): batch_norm_node = helper.make_node( "BatchNormalization", ["x", "s", "bias", "mean", "var"], ["y"], epsilon=1e-2 From 234d8525a55694e8892adf36dbf1d30b1fd6f24b Mon Sep 17 00:00:00 2001 From: viiccwen Date: Tue, 14 Jul 2026 04:05:49 +0000 Subject: [PATCH 2/2] [Lint] Remove stale F821 suppressions Signed-off-by: viiccwen --- python/tvm/s_tir/dlight/gpu/gemv.py | 2 +- python/tvm/s_tir/dlight/gpu/low_batch_gemv.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/python/tvm/s_tir/dlight/gpu/gemv.py b/python/tvm/s_tir/dlight/gpu/gemv.py index c893c3449a95..1c451b964faa 100644 --- a/python/tvm/s_tir/dlight/gpu/gemv.py +++ b/python/tvm/s_tir/dlight/gpu/gemv.py @@ -14,7 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# ruff: noqa: E741, F821 +# ruff: noqa: E741 """A rule for GEMV and DecodeGEMV.""" from functools import reduce diff --git a/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py b/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py index e854ec1e2639..8dedbec4b661 100644 --- a/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py +++ b/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py @@ -14,7 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# ruff: noqa: E741, F821 +# ruff: noqa: E741 """A rule for low-batch GEMM / decode-GEMM using GEMV schedule.""" from functools import reduce