Fix match case in Python-side expr functor (apache#4037)

neo-ai · Oct 18, 2019 · 15b5752 · 15b5752
1 parent f499fd5
commit 15b5752
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/python/tvm/relay/expr_functor.py b/python/tvm/relay/expr_functor.py
@@ -249,7 +249,10 @@ def visit_constructor(self, con):
         return con
 
     def visit_match(self, m):
-        return Match(self.visit(m.data), [Clause(c.lhs, self.visit(c.rhs)) for c in m.clauses])
+        return Match(
+            self.visit(m.data),
+            [Clause(c.lhs, self.visit(c.rhs)) for c in m.clauses],
+            complete=m.complete)
 
     def visit_ref_create(self, r):
         return RefCreate(self.visit(r.value))

diff --git a/tests/python/relay/test_expr_functor.py b/tests/python/relay/test_expr_functor.py
@@ -125,6 +125,16 @@ def test_match():
     p = relay.prelude.Prelude()
     check_visit(p.mod[p.map])
 
+
+def test_match_completeness():
+    p = relay.prelude.Prelude()
+    for completeness in [True, False]:
+        match_expr = relay.adt.Match(p.nil, [], complete=completeness)
+        result_expr = ExprMutator().visit(match_expr)
+        # ensure the mutator doesn't mangle the completeness flag
+        assert result_expr.complete == completeness
+
+
 if __name__ == "__main__":
     test_constant()
     test_tuple()
@@ -139,3 +149,4 @@ def test_match():
     test_ref_write()
     test_memo()
     test_match()
+    test_match_completeness()