[Fix][Relax][ONNX] Fix Range dtype for PrimExpr bounds

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()

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

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

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