pass text printer tests

apache · Mar 6, 2019 · df08054 · df08054
1 parent 7f352c3
commit df08054
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/src/relay/ir/pretty_printer.cc b/src/relay/ir/pretty_printer.cc
@@ -249,6 +249,13 @@ class PrettyPrinter :
       Doc printed_expr;
       if (meta) {
         printed_expr = meta_.GetMetaNode(GetRef<NodeRef>(expr.get()));
+      } else if (GNF_ && gnf && expr.as<LetNode>()) {
+        // wrap GNFed let in brackets
+        printed_expr = Nil();
+        Doc body = Nil();
+        printed_expr << "{";
+        printed_expr << Indent(2, body << "\n" << VisitExpr(expr)) << "\n";
+        printed_expr << "}";
       } else {
         printed_expr = VisitExpr(expr);
       }
@@ -270,6 +277,9 @@ class PrettyPrinter :
         return memo_[expr];
       } else {
         memo_[expr] = printed_expr;
+        if (GNF_ && expr.as<CallNode>()) {
+          printed_expr << PrintOptionalInfo(expr);
+        }
         return printed_expr;
       }
     }

diff --git a/tests/python/relay/test_ir_text_printer.py b/tests/python/relay/test_ir_text_printer.py
@@ -33,8 +33,8 @@ def test_env():
     text = env.astext()
     assert "def @myf" in text
     assert "def @myf" in str(env)
-    assert "%1 = add(%0, %0) # ty=float32" in text
-    assert "%1 = add(%0, %0) # ty=float32" in str(env)
+    assert "add(%0, %0) # ty=float32" in text
+    assert "add(%0, %0) # ty=float32" in str(env)
     show(env.astext(annotate=lambda x: str(x.checked_type.dtype)))
     show(text)