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 <[email protected]> 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 ^^^^^^^^^^^^
