add blank line

apache · Jan 25, 2023 · b9ee9cc · b9ee9cc
1 parent faeaa7f
commit b9ee9cc
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 0 deletions.
diff --git a/src/script/printer/doc_printer/python_doc_printer.cc b/src/script/printer/doc_printer/python_doc_printer.cc
@@ -680,6 +680,10 @@ String DocToPythonScript(Doc doc, const PrinterConfig& cfg) {
     if (cfg->ir_usage.count("relax")) {
       headers.push_back(CommentDoc("from tvm.script import relax as " + cfg->relax_prefix));
     }
+    if (headers.size()) {
+      // Add a blank line between headers and script codes
+      headers.push_back(CommentDoc(""));
+    }
     printer.Append(StmtBlockDoc(headers));
   }
   printer.Append(doc, cfg->path_to_underline);

diff --git a/tests/python/unittest/test_tvmscript_printer_ir.py b/tests/python/unittest/test_tvmscript_printer_ir.py
@@ -39,6 +39,7 @@ def test_ir_module():
         """
 # from tvm.script import ir as I
 # from tvm.script import tir as T
+
 @I.ir_module
 class Module:
     @T.prim_func

diff --git a/tests/python/unittest/test_tvmscript_printer_tir.py b/tests/python/unittest/test_tvmscript_printer_tir.py
@@ -42,6 +42,7 @@ def test_prim_func():
         func,
         expected="""
 # from tvm.script import tir as T
+
 @T.prim_func
 def main(A: T.Buffer((128, 128), "float32"), B: T.Buffer((256, 256), "float32")):
     T.evaluate(0)""",
@@ -64,6 +65,7 @@ def test_prim_func_no_sugar_inlined_buffer():
         func,
         expected="""
 # from tvm.script import tir as T
+
 @T.prim_func
 def main(a: T.handle, B: T.Buffer((256, 256), "float32")):
     A = T.match_buffer(a, (128, 128))
@@ -89,6 +91,7 @@ def test_prim_func_no_sugar_shared_buffer_data():
         func,
         expected="""
 # from tvm.script import tir as T
+
 @T.prim_func
 def main(a: T.handle, b: T.handle):
     A = T.match_buffer(a, (128, 128))
@@ -703,6 +706,7 @@ def block_with_remap_explicitly():
 
     expected_output = """
 # from tvm.script import tir as T
+
 @T.prim_func
 def main():
     # with T.block("root"):
@@ -739,6 +743,7 @@ def root_block_explicitly():
 
     expected_output = """
 # from tvm.script import tir as T
+
 @T.prim_func
 def main():
     # with T.block("root"):