[Bugfix] Fix bb multi-function creation bug (apache#86)

areusch · Aug 14, 2022 · 140ed55 · 140ed55
1 parent c2605ab
commit 140ed55
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/python/tvm/relax/block_builder.py b/python/tvm/relax/block_builder.py
@@ -160,6 +160,7 @@ def _exit_function_scope(self, exc_type, exc_val, exc_tb):
             if not self._is_emit_func_output_called:
                 raise RuntimeError("emit_func_output must be called in a relax function.")
 
+        self._blocks = []
         self._is_emit_func_output_called = False
         BlockBuilder._current = None
 

diff --git a/tests/python/relax/test_blockbuilder.py b/tests/python/relax/test_blockbuilder.py
@@ -159,10 +159,12 @@ def test_multi_functions():
     assert func1.params[0] == x
     assert func1.params[1] == y
     assert func1.name.name_hint == "func1"
+    assert len(func1.body.blocks) == 1
     func2 = mod["func2"]
     assert func2.params[0] == x
     assert func2.params[1] == y
     assert func2.name.name_hint == "func2"
+    assert len(func2.body.blocks) == 1
 
 
 def test_binary_shape_type_deduction():