This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 1ea930e6db [REFACTOR][Relax] Prefer tuple literals in script value
contexts (#20604)
1ea930e6db is described below
commit 1ea930e6db63c9747b9d81a12438864db3e7126e
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Oct 8 07:15:52 2026 -0400
[REFACTOR][Relax] Prefer tuple literals in script value contexts (#20604)
Empty Relax tuples now print as `()`, including void expression
statements. The parser accepts that empty tuple statement form, and
existing tuple bindings and returns with expression fields use ordinary
tuple syntax.
Keep the explicit tuple and primitive-value helpers where conversion
semantics matter, including numeric tuple field types and explicit
scalar dtypes. Document these distinctions and use existing primitive
expressions directly in script examples.
---
python/tvm/relax/expr.py | 13 +++++-
python/tvm/relax/script/ir_builder/ir.py | 10 ++++-
.../tvm/relax/script/ir_builder/parser_protocol.py | 2 +
python/tvm/relax/transform/transform.py | 6 +--
src/relax/script/printer/expr.cc | 1 -
tests/python/disco/test_loader.py | 2 +-
.../relax/script/test_relax_script_basic_usage.py | 48 +++++++++++++++++++++-
.../script/test_relax_script_error_handling.py | 10 +++++
.../relax/script/test_relax_script_printer.py | 4 +-
tests/python/relax/test_analysis.py | 2 +-
tests/python/relax/test_analysis_type_analysis.py | 2 +-
tests/python/relax/test_dataflow_rewriter.py | 8 ++--
tests/python/relax/test_frontend_onnx.py | 4 +-
tests/python/relax/test_inline_functions.py | 2 +-
tests/python/relax/test_transform_fuse_ops.py | 2 +-
tests/python/relax/test_transform_gradient.py | 2 +-
tests/python/relax/test_transform_lambda_lift.py | 4 +-
.../relax/test_transform_lazy_transform_params.py | 28 ++++++-------
.../relax/test_transform_lift_transform_params.py | 4 +-
.../test_transform_merge_composite_functions.py | 4 +-
tests/python/relax/test_transform_normalize.py | 4 +-
.../relax/test_transform_rewrite_cuda_graph.py | 36 ++++++++--------
tests/python/relax/test_vm_callback_function.py | 2 +-
tests/python/script/test_constructor_contracts.py | 48 ++++++++++++++++++++++
24 files changed, 184 insertions(+), 64 deletions(-)
diff --git a/python/tvm/relax/expr.py b/python/tvm/relax/expr.py
index 917dcfd031..245053ef59 100644
--- a/python/tvm/relax/expr.py
+++ b/python/tvm/relax/expr.py
@@ -58,8 +58,17 @@ def prim_value(value: Expr | int | float, dtype: str | None
= None) -> Expr:
Returns
-------
result : Expr
- The converted primitive expression. Existing ``Expr`` inputs are
- returned unchanged.
+ The converted primitive expression. Existing primitive expression
inputs are
+ returned unchanged, even when ``dtype`` is supplied.
+
+ Notes
+ -----
+ Existing primitive expressions can be used directly in TVMScript. Python
+ scalars are context dependent: an unannotated Relax binding or return
+ converts them to tensor constants, while this helper constructs primitive
+ expressions. ``dtype`` selects the type of Python numeric inputs; omitting
+ it uses bool, int64, or float64. The script alias ``R.prim_value`` retains
+ these programmatic conversion semantics.
"""
if tvm.ir.is_prim_expr(value):
return value
diff --git a/python/tvm/relax/script/ir_builder/ir.py
b/python/tvm/relax/script/ir_builder/ir.py
index 42d491a925..71e356df86 100644
--- a/python/tvm/relax/script/ir_builder/ir.py
+++ b/python/tvm/relax/script/ir_builder/ir.py
@@ -231,7 +231,15 @@ def rewriter(rewriter_mod: IRModule | type) ->
PatternMatchingRewriter:
def tuple(*fields: Expr) -> Expr:
- """Create a tuple expression.
+ """Create a tuple expression from already constructed expression fields.
+
+ In script bindings and returns, prefer ordinary Python tuple syntax,
+ including ``()`` and ``(value,)``, when fields are already expressions.
+ This helper materializes an IR tuple outside parser conversion contexts.
+ Numeric fields use core constructor conversion: for example, ``tuple(1)``
+ has an int32 field, whereas a script tuple literal ``(1,)`` has an int64
+ field. Use explicitly typed expressions to make either spelling equivalent.
+
Parameters
----------
*fields : Expr
diff --git a/python/tvm/relax/script/ir_builder/parser_protocol.py
b/python/tvm/relax/script/ir_builder/parser_protocol.py
index f5df2b67e1..d0d429a044 100644
--- a/python/tvm/relax/script/ir_builder/parser_protocol.py
+++ b/python/tvm/relax/script/ir_builder/parser_protocol.py
@@ -419,6 +419,8 @@ def emit_(value: Any, *, span: _Span = None) -> None:
return None
if value is None:
return
+ if isinstance(value, _python.tuple) and not value:
+ value = _relax.Tuple([])
if not isinstance(value, _relax.Expr):
raise TypeError(f"Unsupported expression statement value:
{type(value).__name__}")
result = bind_(_base.at_(span, value), name="_", span=span)
diff --git a/python/tvm/relax/transform/transform.py
b/python/tvm/relax/transform/transform.py
index f64be1e41f..a89c03fe6b 100644
--- a/python/tvm/relax/transform/transform.py
+++ b/python/tvm/relax/transform/transform.py
@@ -576,15 +576,15 @@ def KillAfterLastUse() -> tvm.ir.transform.Pass:
def ComputePrimValue() -> tvm.ir.transform.Pass:
- """Compute all R.prim_value instances
+ """Compute symbolic primitive expressions
While high-level relax can include expressions in terms of its
symbolic variables, these expressions cannot natively be computed
within relax. In order to provide values for symbolic expressions
- (e.g. `R.prim_value(N*N)`, where `N` is a symbolic variable), this
+ (e.g. `N * N`, where `N` is a symbolic variable), this
pass generates a Function in which the expression can be computed.
The relax graph is then updated to include a call to that
- Function, in place of the original `R.prim_value(expr)`.
+ Function, in place of the original primitive expression.
Returns
-------
diff --git a/src/relax/script/printer/expr.cc b/src/relax/script/printer/expr.cc
index 464cb69072..6984e63590 100644
--- a/src/relax/script/printer/expr.cc
+++ b/src/relax/script/printer/expr.cc
@@ -34,7 +34,6 @@ ffi::Optional<ExprDoc> TupleDocTranslate(DocTranslatorObj* d,
ffi::AnyView input
const ffi::Object*) {
const auto* tuple =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
TupleNode>(input);
- if (tuple->fields.empty()) return
NamespaceDoc("relax")->Attr("tuple")->Call({});
ffi::Array<ExprDoc> fields;
for (const Expr& field : tuple->fields)
fields.push_back(d->Translate(field).value());
return TupleDoc(fields);
diff --git a/tests/python/disco/test_loader.py
b/tests/python/disco/test_loader.py
index f661793ffe..6a293aff15 100644
--- a/tests/python/disco/test_loader.py
+++ b/tests/python/disco/test_loader.py
@@ -291,7 +291,7 @@ def test_load_shard_in_relax():
R.shape([1]),
ty_args=R.Tensor((16, 128), "float32"),
)
- lv2 = R.tuple(lv0, lv1)
+ lv2 = (lv0, lv1)
R.output(lv2)
return lv2
diff --git a/tests/python/relax/script/test_relax_script_basic_usage.py
b/tests/python/relax/script/test_relax_script_basic_usage.py
index 9aa534049b..ddc34a887b 100644
--- a/tests/python/relax/script/test_relax_script_basic_usage.py
+++ b/tests/python/relax/script/test_relax_script_basic_usage.py
@@ -307,7 +307,7 @@ def test_tuple_binding():
def test_tuple_get_item():
@R.function
def foo(x: R.Tensor, y: R.Tensor):
- t1 = R.tuple(x, y)
+ t1 = (x, y)
t2 = (x, y)
a = t1[0]
b = R.TupleGetItem(t2, 1)
@@ -868,7 +868,7 @@ def test_tir_expr_as_assert_condition():
def test_empty_tuple():
@R.function
def foo(x: R.Tuple()):
- y: R.Tuple() = R.tuple()
+ y: R.Tuple() = ()
return y
x = relax.Var("x", relax.TupleType([]))
@@ -1488,5 +1488,49 @@ def test_roundtrip_basic_usage(ir_generator):
tvm.ir.assert_structural_equal(original, after_roundtrip,
map_free_vars=True)
+def test_tuple_literals_preserve_expression_structure():
+ @R.function(private=True)
+ def literal(x: T.int64):
+ empty = ()
+ singleton = (x,)
+ nested = (empty, singleton)
+ return nested
+
+ @R.function(private=True)
+ def constructed(x: T.int64):
+ empty = R.tuple()
+ singleton = R.tuple(x)
+ nested = R.tuple(empty, singleton)
+ return nested
+
+ tvm.ir.assert_structural_equal(literal, constructed)
+ for show_all_ty in [False, True]:
+ source = literal.script(show_all_ty=show_all_ty)
+ assert "R.tuple" not in source
+ tvm.ir.assert_structural_equal(
+ literal, tvm.script.from_source(source, extra_vars={"I": I, "R":
R, "T": T})
+ )
+
+
+def test_primitive_scalar_conversion_depends_on_context():
+ @R.function(private=True)
+ def tensor():
+ return 1
+
+ @R.function(private=True)
+ def primitive():
+ return R.prim_value(1)
+
+ @R.function(private=True)
+ def typed():
+ value: T.int32 = 1
+ return value
+
+ assert isinstance(tensor.ret_ty, relax.TensorType)
+ assert primitive.ret_ty.dtype == "int64"
+ assert typed.ret_ty.dtype == "int32"
+ _check(typed)
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/relax/script/test_relax_script_error_handling.py
b/tests/python/relax/script/test_relax_script_error_handling.py
index babe2120ed..578f8f3c0c 100644
--- a/tests/python/relax/script/test_relax_script_error_handling.py
+++ b/tests/python/relax/script/test_relax_script_error_handling.py
@@ -386,6 +386,16 @@ def
test_function_with_non_void_return_type_must_be_assigned():
return x
[email protected]("value", [(1,), ((),), [()]])
+def test_nonempty_host_tuple_statement_is_rejected(value):
+ with pytest.raises(TypeError, match="Unsupported expression statement"):
+
+ @R.function
+ def func():
+ value
+ return ()
+
+
def test_private_function_with_global_symbol_fail():
with pytest.raises(ValueError):
diff --git a/tests/python/relax/script/test_relax_script_printer.py
b/tests/python/relax/script/test_relax_script_printer.py
index ecaedff7d7..86e8aff794 100644
--- a/tests/python/relax/script/test_relax_script_printer.py
+++ b/tests/python/relax/script/test_relax_script_printer.py
@@ -582,7 +582,7 @@ def func() -> T.int64:
@R.function
def float_func() -> T.float32:
- return R.prim_value(T.float32(1.0))
+ return T.float32(1.0)
float_script = float_func.script(verbose_expr=True)
assert "R.prim_value" not in float_script
@@ -605,7 +605,7 @@ def
test_primitive_bindings_roundtrip_without_prim_value_marker():
@R.function
def func(n: T.int64) -> T.int64:
plus_one = n + 1
- alias = R.prim_value(plus_one)
+ alias = plus_one
return alias
for show_all_ty in [False, True]:
diff --git a/tests/python/relax/test_analysis.py
b/tests/python/relax/test_analysis.py
index 59abb7af86..d02c329e65 100644
--- a/tests/python/relax/test_analysis.py
+++ b/tests/python/relax/test_analysis.py
@@ -216,7 +216,7 @@ def
test_binding_block_keep_pure_func_used_only_for_impure():
z = R.call_packed(
"function_maybe_with_side_effects", y, ty_args=(R.Tensor((32, 32),
"int32"))
)
- return R.tuple()
+ return ()
expected = before
diff --git a/tests/python/relax/test_analysis_type_analysis.py
b/tests/python/relax/test_analysis_type_analysis.py
index de64b619a2..0ed2993260 100644
--- a/tests/python/relax/test_analysis_type_analysis.py
+++ b/tests/python/relax/test_analysis_type_analysis.py
@@ -765,7 +765,7 @@ def test_collect_nonnegative_expressions():
C: R.Shape([M, N]),
D: T.int64,
):
- return R.tuple()
+ return ()
M, N = list(func.params[2].ty.values)
diff --git a/tests/python/relax/test_dataflow_rewriter.py
b/tests/python/relax/test_dataflow_rewriter.py
index 6a70501c5f..1480e2b4b3 100644
--- a/tests/python/relax/test_dataflow_rewriter.py
+++ b/tests/python/relax/test_dataflow_rewriter.py
@@ -71,7 +71,7 @@ def test_missing_pattern_raises_error():
class Rewriter:
@R.function
def replacement():
- return R.tuple()
+ return ()
def test_incorrect_function_type_of_pattern_raises_error():
@@ -87,7 +87,7 @@ def test_incorrect_function_type_of_pattern_raises_error():
@R.function
def replacement():
- return R.tuple()
+ return ()
def test_missing_replacement_raises_error():
@@ -99,7 +99,7 @@ def test_missing_replacement_raises_error():
class Rewriter:
@R.function
def pattern():
- return R.tuple()
+ return ()
def test_incorrect_function_type_of_replacement_raises_error():
@@ -111,7 +111,7 @@ def
test_incorrect_function_type_of_replacement_raises_error():
class Rewriter:
@R.function
def pattern():
- return R.tuple()
+ return ()
@Ts.function
def replacement():
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index 5f547cdd9c..e6f43675a9 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -11120,9 +11120,9 @@ def test_sequence_empty():
def main() -> R.Tuple:
R.func_attr({"num_input": 0})
with R.dataflow():
- gv: R.Tuple = R.tuple()
+ gv: R.Tuple = ()
R.output(gv)
- return R.tuple()
+ return ()
tvm.ir.assert_structural_equal(tvm_model, Expected)
diff --git a/tests/python/relax/test_inline_functions.py
b/tests/python/relax/test_inline_functions.py
index 02354e54fe..4dcd5ea059 100644
--- a/tests/python/relax/test_inline_functions.py
+++ b/tests/python/relax/test_inline_functions.py
@@ -76,7 +76,7 @@ def test_ambiguous_function_name():
@R.function
def func():
- return R.tuple()
+ return ()
gvar = tvm.ir.GlobalVar("name")
diff --git a/tests/python/relax/test_transform_fuse_ops.py
b/tests/python/relax/test_transform_fuse_ops.py
index 7a358c69ad..11e811e005 100644
--- a/tests/python/relax/test_transform_fuse_ops.py
+++ b/tests/python/relax/test_transform_fuse_ops.py
@@ -867,7 +867,7 @@ def test_edge_with_call_dps_packed():
b = R.call_tir(cls.exp, (a,), out_ty=R.Tensor((2, 3),
"float32"))
c = R.call_dps_packed("packed_dps", (a,), out_ty=R.Tensor((2,
3), "float32"))
R.output(b, c)
- return R.tuple(b, c)
+ return (b, c)
@Ts.function(private=True)
def exp(A: T.Tensor((2, 3), "float32"), B: T.Tensor((2, 3),
"float32")):
diff --git a/tests/python/relax/test_transform_gradient.py
b/tests/python/relax/test_transform_gradient.py
index eb9104470f..e5451b9ba0 100644
--- a/tests/python/relax/test_transform_gradient.py
+++ b/tests/python/relax/test_transform_gradient.py
@@ -1107,7 +1107,7 @@ def test_report_error():
def main(x: R.Tensor((3, 3), "float32")):
with R.dataflow():
lv1 = R.sum(x)
- gv = R.tuple(lv1, lv1)
+ gv = (lv1, lv1)
R.output(gv)
return gv
diff --git a/tests/python/relax/test_transform_lambda_lift.py
b/tests/python/relax/test_transform_lambda_lift.py
index 30d365b2e7..e6d155ac78 100644
--- a/tests/python/relax/test_transform_lambda_lift.py
+++ b/tests/python/relax/test_transform_lambda_lift.py
@@ -413,7 +413,7 @@ def test_lambda_function_with_same_name_as_global():
@R.function
def main_inner():
- return R.tuple()
+ return ()
@I.ir_module
class Expected:
@@ -433,7 +433,7 @@ def test_lambda_function_with_same_name_as_global():
@R.function
def main_inner():
- return R.tuple()
+ return ()
after = transform.LambdaLift()(Before)
assert_structural_equal(Expected, after)
diff --git a/tests/python/relax/test_transform_lazy_transform_params.py
b/tests/python/relax/test_transform_lazy_transform_params.py
index 8082bd4ef5..ff82c8fb81 100644
--- a/tests/python/relax/test_transform_lazy_transform_params.py
+++ b/tests/python/relax/test_transform_lazy_transform_params.py
@@ -101,7 +101,7 @@ def test_lazy_transform_params():
)
_2: R.Tuple = R.vm.kill_object(lv1_m)
_3: R.Any = R.call_packed("set_item", R.prim_value(1), lv2,
ty_args=(R.Any,))
- gv: R.Tuple = R.tuple()
+ gv: R.Tuple = ()
return gv
after = LazyTransformParams()(Before)
@@ -267,7 +267,7 @@ def test_extra_get_item_params():
_2: R.Tuple = R.vm.kill_object(lv1)
lv3: R.Tensor((16, 3, 3, 3), dtype="float32") = R.add(lv2,
R.const(1, "float32"))
_3: R.Any = R.call_packed("set_item", R.prim_value(1), lv3,
ty_args=(R.Any,))
- gv_1: R.Tuple = R.tuple()
+ gv_1: R.Tuple = ()
return gv_1
after = LazyTransformParams(extra_get_item_params=[relax.Var("loader",
relax.AnyType())])(
@@ -352,7 +352,7 @@ def test_extra_set_item_params():
_2: R.Tuple = R.vm.kill_object(lv1)
lv3: R.Tensor((16, 3, 3, 3), dtype="float32") = R.add(lv2,
R.const(1, "float32"))
_3: R.Any = R.call_packed("set_item", setter, R.prim_value(1),
lv3, ty_args=(R.Any,))
- gv_1: R.Tuple = R.tuple()
+ gv_1: R.Tuple = ()
return gv_1
after = LazyTransformParams(extra_set_item_params=[relax.Var("setter",
relax.AnyType())])(
@@ -379,7 +379,7 @@ def test_extra_set_item_params_with_const_output():
class Expected:
@R.function(pure=False)
def main_transform_params(setter: R.Any) -> R.Tuple:
- output = R.tuple()
+ output = ()
_ = R.call_packed(
"set_item",
setter,
@@ -461,7 +461,7 @@ def test_lazy_transform_params_with_symbolic_vars():
unused_1_ = R.vm.kill_object(param_m)
unused_2_ = R.call_packed("set_item", R.prim_value(0),
transformed, ty_args=(R.Any,))
- output = R.tuple()
+ output = ()
return output
@Ts.function(private=True)
@@ -562,7 +562,7 @@ def test_param_shape_symbolic():
)
_2: R.Tuple = R.vm.kill_object(lv1)
_3: R.Any = R.call_packed("set_item", R.prim_value(1), lv2,
ty_args=(R.Any,))
- gv4: R.Tuple = R.tuple()
+ gv4: R.Tuple = ()
return gv4
after = LazyTransformParams()(Before)
@@ -612,7 +612,7 @@ def test_output_with_use_site():
z = R.call_tir(cls.copy, (y,), out_ty=R.Tensor((),
dtype="float32"))
_1: R.Any = R.call_packed("set_item", R.prim_value(0), y,
ty_args=(R.Any,))
_2: R.Any = R.call_packed("set_item", R.prim_value(1), z,
ty_args=(R.Any,))
- gv: R.Tuple = R.tuple()
+ gv: R.Tuple = ()
return gv
after = LazyTransformParams()(Module)
@@ -716,7 +716,7 @@ def test_duplicate_outputs():
transformed1: R.Tensor((16,), dtype="int32") = R.add(param1,
R.const(2, "int32"))
_ = R.vm.kill_object(param1)
_ = R.call_packed("set_item", R.prim_value(1), transformed1,
ty_args=(R.Any,))
- output = R.tuple()
+ output = ()
return output
after = LazyTransformParams()(Before)
@@ -861,7 +861,7 @@ def test_set_output_callback():
fset_output(R.prim_value(1), C)
D = R.add(C, B)
fset_output(R.prim_value(0), D)
- return R.tuple()
+ return ()
After = relax.transform.LazySetOutput()(Before)
tvm.ir.assert_structural_equal(After, Expected)
@@ -895,7 +895,7 @@ def test_set_output_callback_of_param():
C = R.multiply(A, R.const(2, "float32"))
D = R.add(C, B)
fset_output(R.prim_value(0), D)
- return R.tuple()
+ return ()
After = relax.transform.LazySetOutput()(Before)
tvm.ir.assert_structural_equal(After, Expected)
@@ -930,7 +930,7 @@ def test_set_output_callback_num_input():
C = R.multiply(A, R.const(2, "float32"))
D = R.add(C, B)
fset_output(R.prim_value(0), D)
- return R.tuple()
+ return ()
After = relax.transform.LazySetOutput()(Before)
tvm.ir.assert_structural_equal(After, Expected)
@@ -964,7 +964,7 @@ def test_set_output_callback_with_duplicate_output():
D = R.add(C, B)
fset_output(R.prim_value(0), D)
fset_output(R.prim_value(1), D)
- return R.tuple()
+ return ()
After = relax.transform.LazySetOutput()(Before)
tvm.ir.assert_structural_equal(After, Expected)
@@ -1001,7 +1001,7 @@ def test_set_output_callback_with_inline_const():
fset_output(R.prim_value(1), D)
fset_output(R.prim_value(2), R.prim_value(42))
fset_output(R.prim_value(3), R.const(17.5, "float16"))
- return R.tuple()
+ return ()
After = relax.transform.LazySetOutput()(Before)
tvm.ir.assert_structural_equal(After, Expected)
@@ -1029,7 +1029,7 @@ def test_set_output_callback_with_non_tuple_output():
C = R.multiply(A, R.const(2, "float32"))
D = R.add(C, B)
fset_output(R.prim_value(0), D)
- return R.tuple()
+ return ()
After = relax.transform.LazySetOutput()(Before)
tvm.ir.assert_structural_equal(After, Expected)
diff --git a/tests/python/relax/test_transform_lift_transform_params.py
b/tests/python/relax/test_transform_lift_transform_params.py
index 92882bf3d1..10ed9a6221 100644
--- a/tests/python/relax/test_transform_lift_transform_params.py
+++ b/tests/python/relax/test_transform_lift_transform_params.py
@@ -1420,7 +1420,7 @@ def test_symbolic_var_1():
R.func_attr({"num_input": 0})
# All instance of the empty tuple are normalized to be
# in-line.
- return R.tuple()
+ return ()
@R.function
def main(shape: R.Shape([n])) -> R.Shape([n]):
@@ -1483,7 +1483,7 @@ def test_symbolic_var_2():
@R.function
def main_transform_params(params: R.Tuple) -> R.Tuple:
R.func_attr({"num_input": 0})
- return R.tuple()
+ return ()
@R.function
def main(shape: R.Shape([n_main])) -> R.Shape([n_main]):
diff --git a/tests/python/relax/test_transform_merge_composite_functions.py
b/tests/python/relax/test_transform_merge_composite_functions.py
index b679492c58..c619ef730a 100644
--- a/tests/python/relax/test_transform_merge_composite_functions.py
+++ b/tests/python/relax/test_transform_merge_composite_functions.py
@@ -1259,7 +1259,7 @@ def test_tuple_projection_merging():
parts = cls.split(x)
left = parts[0]
right = parts[1]
- repacked = R.tuple(left, right)
+ repacked = (left, right)
reordered_left = repacked[0]
reordered_right = repacked[1]
difference = cls.subtract(reordered_right, reordered_left)
@@ -1375,7 +1375,7 @@ def test_tuple_projection_rejects_escaping_tuple():
alias = parts
left = alias[0]
out = cls.relu(left)
- result = R.tuple(alias, out)
+ result = (alias, out)
R.output(result)
return result
diff --git a/tests/python/relax/test_transform_normalize.py
b/tests/python/relax/test_transform_normalize.py
index 26cdcb78c9..342649804d 100644
--- a/tests/python/relax/test_transform_normalize.py
+++ b/tests/python/relax/test_transform_normalize.py
@@ -566,7 +566,7 @@ def test_remove_usage_of_void_type_variables():
`void_var` may require a binding.
This is avoided by normalizing all usage of a void-type
- variable with an in-line `R.tuple()`.
+ variable with an in-line empty tuple.
"""
x = relax.Var("x", R.Tuple([]))
bindings = [
@@ -580,7 +580,7 @@ def test_remove_usage_of_void_type_variables():
@R.function(private=True, pure=False)
def expected():
x = R.assert_op(R.const(True, "bool"))
- return R.tuple()
+ return ()
if __name__ == "__main__":
diff --git a/tests/python/relax/test_transform_rewrite_cuda_graph.py
b/tests/python/relax/test_transform_rewrite_cuda_graph.py
index 4737c127ee..4e9b85f66c 100644
--- a/tests/python/relax/test_transform_rewrite_cuda_graph.py
+++ b/tests/python/relax/test_transform_rewrite_cuda_graph.py
@@ -696,7 +696,7 @@ def test_transform_is_no_op_when_disabled():
def main():
storage = R.memory.alloc_storage(R.shape([8]), 0, "global",
"float32")
alloc3 = R.memory.alloc_tensor(storage, 0, R.shape([8]), "float32")
- return R.tuple()
+ return ()
with tvm.transform.PassContext(config={"relax.backend.use_cuda_graph":
True}):
AfterWhenEnabled = relax.transform.RewriteCUDAGraph()(Before)
@@ -715,7 +715,7 @@ def test_static_args():
storage0 = R.memory.alloc_storage(R.shape([8]), 0, "global",
"float32")
alloc0 = R.memory.alloc_tensor(storage0, 0, R.shape([8]),
"float32")
_ = R.call_packed("dummy_func", alloc0, R.dtype("float32"),
R.str("string"))
- return R.tuple()
+ return ()
@I.ir_module
class Expected:
@@ -732,7 +732,7 @@ def test_static_args():
def main_cuda_graph_capture(alloc0: R.Tensor((8,), dtype="float32"))
-> R.Tuple:
R.func_attr({"relax.force_pure": True})
_: R.Any = R.call_packed("dummy_func", alloc0, R.dtype("float32"),
R.str("string"))
- gv: R.Tuple = R.tuple()
+ gv: R.Tuple = ()
return gv
@R.function(pure=False)
@@ -752,7 +752,7 @@ def test_static_args():
(cls.main_cuda_graph_capture, (alloc0,), R.prim_value(0)),
ty_args=(R.Tuple,),
)
- return R.tuple()
+ return ()
mod = relax.transform.RewriteCUDAGraph()(Before)
tvm.ir.assert_structural_equal(mod, Expected)
@@ -833,8 +833,8 @@ def test_dynamic_capture():
R.func_attr({"relax.force_pure": True})
cls = Expected
R.call_tir_packed(cls.add_one, (alloc1, alloc2))
- gv = R.tuple()
- return R.tuple()
+ gv = ()
+ return ()
@R.function
def main(x: R.Tensor((m_main,), dtype="float32")) ->
R.Tensor((m_main,), dtype="float32"):
@@ -889,7 +889,7 @@ def test_merge_alloc_funcs():
alloc2 = R.memory.alloc_tensor(storage2, 0, R.shape([256]),
"float32")
alloc3 = R.memory.alloc_tensor(storage3, 0, R.shape([512]),
"float32")
R.call_packed("dummy", alloc1, alloc2, alloc3, ty_args=(R.Tuple,))
- return R.tuple()
+ return ()
@R.function
def func2():
@@ -903,7 +903,7 @@ def test_merge_alloc_funcs():
alloc3 = R.memory.alloc_tensor(storage3, 0, R.shape([1024]),
"float32")
alloc4 = R.memory.alloc_tensor(storage4, 0, R.shape([512]),
"float32")
R.call_packed("dummy", alloc1, alloc2, alloc3, alloc4,
ty_args=(R.Tuple,))
- return R.tuple()
+ return ()
@I.ir_module
class Expected:
@@ -956,7 +956,7 @@ def test_merge_alloc_funcs():
(cls.func1_cuda_graph_capture, (alloc1, alloc2, alloc3),
R.prim_value(0)),
ty_args=(R.Tuple,),
)
- return R.tuple()
+ return ()
@R.function(private=True)
def func1_cuda_graph_capture(
@@ -966,8 +966,8 @@ def test_merge_alloc_funcs():
) -> R.Tuple:
R.func_attr({"relax.force_pure": True})
R.call_packed("dummy", alloc1, alloc2, alloc3, ty_args=(R.Tuple,))
- R.tuple()
- return R.tuple()
+ ()
+ return ()
@R.function
def func2() -> R.Tuple:
@@ -999,7 +999,7 @@ def test_merge_alloc_funcs():
(cls.func2_cuda_graph_capture, (alloc1, alloc2, alloc3,
alloc4), R.prim_value(1)),
ty_args=(R.Tuple,),
)
- return R.tuple()
+ return ()
@R.function(private=True)
def func2_cuda_graph_capture(
@@ -1010,8 +1010,8 @@ def test_merge_alloc_funcs():
) -> R.Tuple:
R.func_attr({"relax.force_pure": True})
R.call_packed("dummy", alloc1, alloc2, alloc3, alloc4,
ty_args=(R.Tuple,))
- R.tuple()
- return R.tuple()
+ ()
+ return ()
After = relax.transform.RewriteCUDAGraph()(Before)
tvm.ir.assert_structural_equal(After, Expected)
@@ -1055,8 +1055,8 @@ def test_disable_capture_output():
) -> R.Tuple:
R.func_attr({"relax.force_pure": True})
R.call_packed("dummy", alloc1, alloc2, ty_args=(R.Tuple,))
- R.tuple()
- return R.tuple()
+ ()
+ return ()
@R.function
def main(x: R.Tensor((8,), dtype="float32")) -> R.Tuple(R.Tensor((8,),
dtype="float32")):
@@ -1141,8 +1141,8 @@ def test_static_input_with_symbolic_shape():
) -> R.Tuple:
R.func_attr({"relax.force_pure": True})
R.call_packed("dummy", alloc1, w, alloc2, ty_args=(R.Tuple,))
- R.tuple()
- return R.tuple()
+ ()
+ return ()
@R.function
def main(x: R.Tensor((8,), dtype="float16"), w: R.Tensor((m_main,)))
-> R.Tuple(
diff --git a/tests/python/relax/test_vm_callback_function.py
b/tests/python/relax/test_vm_callback_function.py
index cc3d0d8be5..276c3af01d 100644
--- a/tests/python/relax/test_vm_callback_function.py
+++ b/tests/python/relax/test_vm_callback_function.py
@@ -37,7 +37,7 @@ def test_pass_tensor_to_function():
):
B = R.multiply(A, R.const(2))
_ = callback(B)
- return R.tuple()
+ return ()
ex = tvm.relax.build(
tvm.IRModule.from_expr(relax_func),
diff --git a/tests/python/script/test_constructor_contracts.py
b/tests/python/script/test_constructor_contracts.py
index 35537b0311..49482607c0 100644
--- a/tests/python/script/test_constructor_contracts.py
+++ b/tests/python/script/test_constructor_contracts.py
@@ -110,5 +110,53 @@ def test_normal_constructors_in_parsed_function():
ir.assert_structural_equal(identity.params[0].ty, relax.TensorType([2],
"float32"))
[email protected]("dtype,value", [("int32", 3), ("float32", 1.5),
("bool", True)])
+def test_prim_value_explicit_dtype_and_identity(dtype, value):
+ converted = R.prim_value(value, dtype=dtype)
+ assert converted.ty.dtype == dtype
+ assert R.prim_value(converted, dtype="float64").same_as(converted)
+
+
+def test_prim_value_boolean_and_symbolic_conversion():
+ assert R.prim_value(True).ty.dtype == "bool"
+ assert relax.utils.convert_to_expr(True).ty.dtype == "int64"
+ n = tirx.Var("n", "int64")
+ value = n + 1
+ assert R.prim_value(value).same_as(value)
+
+
[email protected]("value", [None, {}])
+def test_expression_helpers_reject_non_expression_inputs(value):
+ with pytest.raises((TypeError, ValueError)):
+ R.prim_value(value)
+ with pytest.raises((TypeError, ValueError)):
+ R.tuple(value)
+
+
+def test_tuple_constructor_materializes_raw_call_arguments():
+ value = tirx.IntImm("int64", 1)
+ nested = R.tuple(R.tuple(), R.tuple(value))
+ ir.assert_structural_equal(nested, ir.Tuple([ir.Tuple([]),
ir.Tuple([value])]))
+ call = ir.Call.unchecked("relax.add", [nested], attrs={"tag": 1},
ty=ir.Type.missing())
+ assert call.args[0].same_as(nested)
+ from_python = ir.Call.unchecked(
+ "relax.add", [((), (value,))], attrs={"tag": 1}, ty=ir.Type.missing()
+ )
+ ir.assert_structural_equal(call, from_python)
+ source = call.script()
+ ir.assert_structural_equal(call, eval(source, {"I": I, "R": R, "T": T}))
+
+
[email protected](
+ "value,constructor_dtype,parser_dtype",
+ [(1, "int32", "int64"), (1.0, "float32", "float64"), (True, "bool",
"int64")],
+)
+def test_tuple_numeric_conversion_depends_on_context(value, constructor_dtype,
parser_dtype):
+ constructed = R.tuple(value)
+ converted = relax.utils.convert_to_expr((value,))
+ assert constructed.fields[0].ty.dtype == constructor_dtype
+ assert converted.fields[0].ty.dtype == parser_dtype
+
+
if __name__ == "__main__":
tvm.testing.main()
