[Relax][Frontend][ONNX] Support dynamic Range bounds

python/tvm/relax/frontend/onnx/onnx_frontend.py

@@ -3724,24 +3724,62 @@ 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
-
-        if isinstance(start, relax.Constant):
-            start = start.data.numpy().tolist()
+        out_dtype = str(start.ty.dtype)
+
+        def get_scalar_value(x):
+            if isinstance(x, relax.Constant):
+                value = x.data.numpy()
+                if value.size != 1:
+                    raise ValueError("Range scalar input must have exactly one element.")
+                return value.item()
+            return x
 
-        if isinstance(limit, relax.Constant):
-            limit = limit.data.numpy().tolist()
+        start = get_scalar_value(start)
+        limit = get_scalar_value(limit)
+        delta = get_scalar_value(delta)
 
-        assert isinstance(delta, relax.Constant), "Constant delta required for Range."
-        step = delta.data.numpy().tolist()
+        has_tensor_expr = any(
+            isinstance(x, relax.Expr) and not tvm.ir.is_prim_expr(x) for x in [start, limit, delta]
+        )
+        has_prim_expr = any(tvm.ir.is_prim_expr(x) for x in [start, limit, delta])
 
-        # If all inputs are constant, compute directly.
-        if isinstance(start, int) and isinstance(limit, int):
-            out_range = _np.arange(start=start, stop=limit, step=step)
+        if not has_tensor_expr and not has_prim_expr:
+            out_range = _np.arange(start=start, stop=limit, step=delta)
             return relax.const(out_range, out_dtype)
 
-        # Otherwise compute in graph.
-        return relax.op.arange(start, limit, step, out_dtype)
+        if not has_tensor_expr:
+            return relax.op.arange(start, limit, delta, out_dtype)
+
+        work_dtype = out_dtype if out_dtype.startswith("float") else "int64"
+
+        def scalar_expr(x):
+            if isinstance(x, relax.Expr):
+                return bb.normalize(relax.op.astype(x, work_dtype))
+            return relax.const(x, work_dtype)
+
+        start_expr = scalar_expr(start)
+        limit_expr = scalar_expr(limit)
+        delta_expr = scalar_expr(delta)
+
+        if work_dtype.startswith("float"):
+            count = relax.op.ceil(
+                relax.op.divide(relax.op.subtract(limit_expr, start_expr), delta_expr)
+            )
+        else:
+            count = relax.op.negative(
+                relax.op.floor_divide(relax.op.subtract(start_expr, limit_expr), delta_expr)
+            )
+
+        count = bb.normalize(relax.op.maximum(count, relax.const(0, work_dtype)))
+        count = bb.normalize(relax.op.astype(count, "int64"))
+        count = bb.normalize(relax.op.reshape(count, (1,)))
+        range_len = _tensor_to_shape_expr(bb, count, 1, "range_len").values[0]
+
+        positions = bb.normalize(
+            relax.op.astype(relax.op.arange(0, range_len, 1, "int64"), work_dtype)
+        )
+        output = relax.op.add(relax.op.multiply(positions, delta_expr), start_expr)
+        return output if work_dtype == out_dtype else relax.op.astype(output, out_dtype)
 
 
 class InstanceNormalization(OnnxOpConverter):

tests/python/relax/test_frontend_onnx.py

@@ -8216,6 +8216,84 @@ def main(
     tvm.ir.assert_structural_equal(tvm_model, Expected)
 
 
+@pytest.mark.parametrize(
+    "start, limit, delta, tensor_dtype, np_dtype",
+    [
+        (0, 6, 2, TensorProto.INT64, np.int64),
+        (8, 0, -2, TensorProto.INT64, np.int64),
+        (5, 1, 1, TensorProto.INT64, np.int64),
+        (0, 7, 2, TensorProto.INT32, np.int32),
+        (0.0, 1.0, 0.25, TensorProto.FLOAT, np.float32),
+        (1.0, -1.0, -0.5, TensorProto.FLOAT, np.float32),
+    ],
+)
+def test_range_dynamic_scalar_inputs(start, limit, delta, tensor_dtype, np_dtype):
+    range_node = helper.make_node(
+        "Range",
+        ["start", "limit", "delta"],
+        ["output"],
+    )
+
+    graph = helper.make_graph(
+        [range_node],
+        "range_dynamic_scalar_inputs_test",
+        inputs=[
+            helper.make_tensor_value_info("start", tensor_dtype, []),
+            helper.make_tensor_value_info("limit", tensor_dtype, []),
+            helper.make_tensor_value_info("delta", tensor_dtype, []),
+        ],
+        outputs=[
+            helper.make_tensor_value_info("output", tensor_dtype, ["range_len"]),
+        ],
+    )
+
+    model = helper.make_model(graph, producer_name="range_dynamic_scalar_inputs_test")
+    check_correctness(
+        model,
+        inputs={
+            "start": np.array(start, dtype=np_dtype),
+            "limit": np.array(limit, dtype=np_dtype),
+            "delta": np.array(delta, dtype=np_dtype),
+        },
+        opset=12,
+        check_dtypes=True,
+    )
+
+
+def test_range_symbolic_primexpr_limit():
+    shape = helper.make_node("Shape", ["x"], ["x_shape"])
+    axis = make_constant_node("axis", TensorProto.INT64, [], [1])
+    gather = helper.make_node("Gather", ["x_shape", "axis"], ["limit_int"])
+    cast = helper.make_node("Cast", ["limit_int"], ["limit"], to=TensorProto.FLOAT)
+    start = make_constant_node("start", TensorProto.FLOAT, [], [0.0])
+    delta = make_constant_node("delta", TensorProto.FLOAT, [], [1.0])
+    range_node = helper.make_node(
+        "Range",
+        ["start", "limit", "delta"],
+        ["output"],
+    )
+
+    graph = helper.make_graph(
+        [shape, axis, gather, cast, start, delta, range_node],
+        "range_symbolic_primexpr_limit_test",
+        inputs=[
+            helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, "range_len"]),
+        ],
+        outputs=[
+            helper.make_tensor_value_info("output", TensorProto.FLOAT, ["range_len"]),
+        ],
+    )
+
+    model = helper.make_model(
+        graph,
+        producer_name="range_symbolic_primexpr_limit_test",
+        opset_imports=[helper.make_opsetid("", 17)],
+    )
+    model.ir_version = 8
+
+    from_onnx(model, opset=17, keep_params_in_input=True)
+
+
 def test_batch_norm():
     batch_norm_node = helper.make_node(
         "BatchNormalization", ["x", "s", "bias", "mean", "var"], ["y"], epsilon=1e-2