You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/10/01 22:58:04 UTC
(tvm) 08/08: Keep complete output expectations aligned with renderer headers
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch tqchen/printer-refinement
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 6927b91cf913ce9d89637abc7d7fbc7d8b46dd94
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 1 22:54:39 2026 +0000
Keep complete output expectations aligned with renderer headers
---
.../relax/script/test_relax_script_printer.py | 2 +-
.../s_tir/script/test_s_tir_script_printer.py | 38 ++++++++++++----------
.../python/tirx/script/test_tirx_script_printer.py | 1 -
.../script/test_tirx_script_printer_underlining.py | 7 ++--
4 files changed, 23 insertions(+), 25 deletions(-)
diff --git a/tests/python/relax/script/test_relax_script_printer.py b/tests/python/relax/script/test_relax_script_printer.py
index a21098dc2b..08b12bc2e8 100644
--- a/tests/python/relax/script/test_relax_script_printer.py
+++ b/tests/python/relax/script/test_relax_script_printer.py
@@ -1179,7 +1179,7 @@ from __future__ import annotations
@R.function
def func(A: R.Tensor((10, 20), dtype="float32"), B: R.Tensor(dtype="float32", ndim=2)) -> R.Tensor((10, 20), dtype="float32"):
B2 = R.match_cast(B, R.Tensor((10, 20), dtype="float32"))
- C = R.add(A, B2)
+ C: R.Tensor((10, 20), dtype="float32") = R.add(A, B2)
D = C
E: R.Tensor((10, 20), dtype="float32") = R.add(D, B)
return E""",
diff --git a/tests/python/s_tir/script/test_s_tir_script_printer.py b/tests/python/s_tir/script/test_s_tir_script_printer.py
index bea1c3b4fb..5d9c9b07a8 100644
--- a/tests/python/s_tir/script/test_s_tir_script_printer.py
+++ b/tests/python/s_tir/script/test_s_tir_script_printer.py
@@ -129,9 +129,10 @@ def test_prim_func():
_assert_print(
func,
expected="""
-# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
+from __future__ import annotations
+
# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@Ts.prim_func
def main(A: T.Buffer((128, 128), "float32", layout="default"), B: T.Buffer((256, 256), "float32", layout="default")):
@@ -154,9 +155,10 @@ def test_prim_func_buffer_data_use():
_assert_print(
func,
expected="""
-# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
+from __future__ import annotations
+
# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@Ts.prim_func
def main(A: T.Buffer((128, 128), "float32", layout="default"), B: T.Buffer((256, 256), "float32", layout="default")):
@@ -181,9 +183,10 @@ def test_prim_func_buffer_data_argument_is_scope_hint():
_assert_print(
func,
expected="""
-# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
+from __future__ import annotations
+
# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@Ts.prim_func
def main(A: T.Buffer((128, 128), "float32", layout="default"), B: T.Buffer((256, 256), "float32", layout="default")):
@@ -389,9 +392,10 @@ def test_private_primfunc():
_assert_print(
func,
expected="""
-# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
+from __future__ import annotations
+
# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@Ts.prim_func(private=True)
def main(A: T.Buffer((128, 128), "float32", layout="default"), B: T.Buffer((256, 256), "float32", layout="default")):
@@ -464,9 +468,8 @@ def test_custom_float_types(dtype):
TB.evaluate(getattr(TB, dtype)(0.0))
expected_output = f"""
-# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@Ts.prim_func
def func():
@@ -2082,8 +2085,8 @@ def test_annotation_multi_access_paths():
)
assert (
result
- == """# from tvm.script import tirx as T
-# from tvm.script import s_tir as Ts
+ == """# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@Ts.prim_func
def main():
@@ -2109,8 +2112,8 @@ def test_annotate_from_multi_obj():
)
assert (
result
- == """# from tvm.script import tirx as T
-# from tvm.script import s_tir as Ts
+ == """# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@Ts.prim_func
def main():
@@ -2141,8 +2144,8 @@ def test_disable_concise_scoping_when_scope_annotated():
)
assert (
result
- == """# from tvm.script import tirx as T
-# from tvm.script import s_tir as Ts
+ == """# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@Ts.prim_func
def main():
@@ -2169,9 +2172,8 @@ def test_ir_module():
mod,
"""
# from tvm.script import ir as I
-# from tvm.script import tirx as T
-# from tvm.tirx.layout import Axis
# from tvm.script import s_tir as Ts
+# from tvm.script import tirx as T
@I.ir_module
class Module:
diff --git a/tests/python/tirx/script/test_tirx_script_printer.py b/tests/python/tirx/script/test_tirx_script_printer.py
index 6e455f934e..b55e962323 100644
--- a/tests/python/tirx/script/test_tirx_script_printer.py
+++ b/tests/python/tirx/script/test_tirx_script_printer.py
@@ -674,7 +674,6 @@ def test_print_kwargs_schedule_op_full_code():
expected = (
"# from tvm.script import tirx as T\n"
- "# from tvm.tirx.layout import Axis\n"
"\n"
"@T.prim_func\n"
"def test():\n"
diff --git a/tests/python/tirx/script/test_tirx_script_printer_underlining.py b/tests/python/tirx/script/test_tirx_script_printer_underlining.py
index 854ac67158..6d8706d763 100644
--- a/tests/python/tirx/script/test_tirx_script_printer_underlining.py
+++ b/tests/python/tirx/script/test_tirx_script_printer_underlining.py
@@ -39,8 +39,9 @@ def test_underline_from_obj():
)
assert result == format_script(
"""
+ from __future__ import annotations
+
# from tvm.script import tirx as T
- # from tvm.tirx.layout import Axis
@T.prim_func
def main(a: T.int32, b: T.int32):
@@ -102,7 +103,6 @@ def test_underline_from_multi_obj():
assert result == format_script(
"""
# from tvm.script import tirx as T
- # from tvm.tirx.layout import Axis
@T.prim_func
def main():
@@ -136,7 +136,6 @@ def test_underline_func():
assert result == format_script(
"""
# from tvm.script import tirx as T
- # from tvm.tirx.layout import Axis
@T.prim_func
^^^^^^^^^^^^
@@ -165,7 +164,6 @@ def test_underline_func_in_irmodule():
"""
# from tvm.script import ir as I
# from tvm.script import tirx as T
- # from tvm.tirx.layout import Axis
@I.ir_module
class Module:
@@ -196,7 +194,6 @@ def test_underline_irmodule():
"""
# from tvm.script import ir as I
# from tvm.script import tirx as T
- # from tvm.tirx.layout import Axis
@I.ir_module
^^^^^^^^^^^^