This is an automated email from the ASF dual-hosted git repository.
yaxingcai 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 5456fae776 [TVMScript,Fix] Fix findsource when classes are indented
(#13924)
5456fae776 is described below
commit 5456fae77655fadf47f61a94e28027c6d39b1e7e
Author: Tristan Konolige <[email protected]>
AuthorDate: Mon Feb 6 20:41:31 2023 -0800
[TVMScript,Fix] Fix findsource when classes are indented (#13924)
Leaving class definitions was not correctly handled when recreating
scoping information. The fix correctly pops scope whenever the
indentation level becomes less than the current scope.
---
python/tvm/script/parser/core/diagnostics.py | 5 +++--
tests/python/unittest/test_tvmscript_parser_source.py | 15 +++++++++++++++
2 files changed, 18 insertions(+), 2 deletions(-)
diff --git a/python/tvm/script/parser/core/diagnostics.py
b/python/tvm/script/parser/core/diagnostics.py
index d673e0eb13..ad7ae50347 100644
--- a/python/tvm/script/parser/core/diagnostics.py
+++ b/python/tvm/script/parser/core/diagnostics.py
@@ -163,9 +163,10 @@ def findsource(obj):
name = tokens[1].split(":")[0].split("(")[0] + "<locals>"
elif tokens[0] == "class":
name = tokens[1].split(":")[0].split("(")[0]
+ # pop scope if we are less indented
+ while scope_stack and indent_info[scope_stack[-1]] >= indent:
+ scope_stack.pop()
if name:
- while scope_stack and indent_info[scope_stack[-1]] >= indent:
- scope_stack.pop()
scope_stack.append(name)
indent_info[name] = indent
if scope_stack == qual_names:
diff --git a/tests/python/unittest/test_tvmscript_parser_source.py
b/tests/python/unittest/test_tvmscript_parser_source.py
index f5dc17fdfe..359583c1aa 100644
--- a/tests/python/unittest/test_tvmscript_parser_source.py
+++ b/tests/python/unittest/test_tvmscript_parser_source.py
@@ -82,5 +82,20 @@ def test_source_ast():
assert isinstance(for_block, doc.With) and len(for_block.body) == 2
+def test_nesting_parsing():
+ class dummy:
+ pass
+
+ for i in range(1):
+
+ @tvm.script.ir_module
+ class Module:
+ @T.prim_func
+ def impl(
+ A: T.Buffer[(12, 196, 64), "float32"],
+ ) -> None:
+ T.evaluate(0)
+
+
if __name__ == "__main__":
tvm.testing.main()
