From 3aecc7ae9dc3c0c91f456c1395a03ea8bb2f40b1 Mon Sep 17 00:00:00 2001 From: Tai An Date: Sun, 6 Sep 2026 18:18:21 -0700 Subject: [PATCH] [S-TIR][Test] Use tvm.testing.main() so schedule tests can run standalone test_tir_schedule_storage_align.py and test_tir_schedule_read_write_at.py both call `tvm.testing.*` without importing `tvm.testing`. That only works under pytest, where the plugin has already imported the submodule; running either file directly raises AttributeError: module 'tvm' has no attribute 'testing' test_tir_schedule_storage_align.py has a second problem: its hand-written `__main__` block calls test_storage_align() but that test takes the parametrized fixture `use_block_name` (`tvm.testing.parameter(by_dict={"block_obj": False, "block_name": True})`), so the call raises TypeError: test_storage_align() missing 1 required positional argument: 'use_block_name' 36 of the 39 files in tests/python/s_tir/schedule already end with `tvm.testing.main()`, which runs the whole module including parametrized tests. Adopt that idiom here and add the missing imports. Signed-off-by: Tai An --- .../s_tir/schedule/test_tir_schedule_read_write_at.py | 1 + .../s_tir/schedule/test_tir_schedule_storage_align.py | 10 ++-------- 2 files changed, 3 insertions(+), 8 deletions(-) diff --git a/tests/python/s_tir/schedule/test_tir_schedule_read_write_at.py b/tests/python/s_tir/schedule/test_tir_schedule_read_write_at.py index 85e6a7a0e0ae..0db98fac2919 100644 --- a/tests/python/s_tir/schedule/test_tir_schedule_read_write_at.py +++ b/tests/python/s_tir/schedule/test_tir_schedule_read_write_at.py @@ -37,6 +37,7 @@ import pytest import tvm +import tvm.testing from tvm import tirx from tvm.s_tir.schedule.testing import ( assert_structural_equal_ignore_global_symbol, diff --git a/tests/python/s_tir/schedule/test_tir_schedule_storage_align.py b/tests/python/s_tir/schedule/test_tir_schedule_storage_align.py index 2280292ec1c9..fd64f2ad8a22 100644 --- a/tests/python/s_tir/schedule/test_tir_schedule_storage_align.py +++ b/tests/python/s_tir/schedule/test_tir_schedule_storage_align.py @@ -19,6 +19,7 @@ import pytest import tvm +import tvm.testing from tvm import tirx from tvm.s_tir.schedule.testing import ( assert_structural_equal_ignore_global_symbol, @@ -173,11 +174,4 @@ def test_storage_align_invalid_annotation(): if __name__ == "__main__": - test_storage_align() - test_storage_align_update() - test_storage_align_invalid_factor1() - test_storage_align_invalid_factor2() - test_storage_align_invalid_buffer() - test_storage_align_invalid_buffer_index() - test_storage_align_invalid_axis() - test_storage_align_invalid_annotation() + tvm.testing.main()