[Relay] Add printing for ADT Type (#3030)

* Update pretty_printer.cc * Update pretty_printer.cc
apache · Apr 17, 2019 · 4ab97df · 4ab97df
1 parent 12257dd
commit 4ab97df
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/src/relay/ir/pretty_printer.cc b/src/relay/ir/pretty_printer.cc
@@ -592,6 +592,22 @@ class PrettyPrinter :
     return AllocTypeVar(GetRef<TypeVar>(node));
   }
 
+  Doc VisitType_(const GlobalTypeVarNode* node) final {
+    return Doc(node->var->name_hint);
+  }
+
+  Doc VisitType_(const TypeCallNode* node) final {
+    Doc doc = PrintType(node->func, false);
+    std::vector<Doc> args;
+    for (const Type& t : node->args) {
+      args.push_back(PrintType(t, false));
+    }
+    doc << "[";
+    doc << PrintVec(args);
+    doc << "]";
+    return doc;
+  }
+
   Doc VisitType_(const TensorTypeNode* node) final {
     // scalar type
     if (node->shape.size() == 0) {