From dbeba1f96ccc73c94c992bf30ac71a7aa55c1505 Mon Sep 17 00:00:00 2001
From: Altan Haan <ahaan@octoml.ai>
Date: Tue, 4 May 2021 23:32:46 +0000
Subject: [PATCH 1/2] fix parsing hierarchical attibute names

---
 src/parser/parser.cc                 | 31 ++++++++++++++++------------
 tests/python/relay/test_ir_parser.py |  7 +++++++
 2 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/src/parser/parser.cc b/src/parser/parser.cc
index b72a632635d9..d234c329a9f1 100644
--- a/src/parser/parser.cc
+++ b/src/parser/parser.cc
@@ -1371,7 +1371,8 @@ class Parser {
     DLOG(INFO) << "Parser::ParseAttrs";
     Map<String, ObjectRef> kwargs;
     while (Peek()->token_type == TokenType::kIdentifier) {
-      auto key = Match(TokenType::kIdentifier).ToString();
+      auto key = GetHierarchicalName(ParseHierarchicalName().data);
+      // auto key = Match(TokenType::kIdentifier).ToString();
       Match(TokenType::kEqual);
       // TOOD(@jroesch): syntactically what do we allow to appear in attribute right hand side.
       auto value = ParseAttributeValue();
@@ -1545,18 +1546,7 @@ class Parser {
             auto spanned_idents = ParseHierarchicalName();
             auto idents = spanned_idents.data;
             auto span = spanned_idents.span;
-            ICHECK_NE(idents.size(), 0);
-            std::stringstream op_name;
-            int i = 0;
-            int periods = idents.size() - 1;
-            for (auto ident : idents) {
-              op_name << ident;
-              if (i < periods) {
-                op_name << ".";
-                i++;
-              }
-            }
-            return GetOp(op_name.str(), span);
+            return GetOp(GetHierarchicalName(idents), span);
           }
         }
         case TokenType::kGraph: {
@@ -1696,6 +1686,21 @@ class Parser {
     return Spanned<Array<String>>(idents, span);
   }
 
+  std::string GetHierarchicalName(Array<String> idents) {
+    ICHECK_NE(idents.size(), 0);
+    std::stringstream hierarchical_name;
+    int i = 0;
+    int periods = idents.size() - 1;
+    for (auto ident : idents) {
+      hierarchical_name << ident;
+      if (i < periods) {
+        hierarchical_name << ".";
+        i++;
+      }
+    }
+    return hierarchical_name.str();
+  }
+
   /*! \brief Parse a shape. */
   Array<tvm::PrimExpr> ParseShape() {
     auto dims = ParseSequence<tvm::PrimExpr>(
diff --git a/tests/python/relay/test_ir_parser.py b/tests/python/relay/test_ir_parser.py
index 8b6b39e3df15..04cc2c0e79e4 100644
--- a/tests/python/relay/test_ir_parser.py
+++ b/tests/python/relay/test_ir_parser.py
@@ -959,6 +959,13 @@ def test_tokenize_inf():
     mod = relay.transform.AnnotateSpans()(mod)
 
 
+def test_func_attrs():
+    attrs = tvm.ir.make_node("DictAttrs", **{"Primitive": 1, "relay.reshape_only": 1})
+    x = relay.var("x", shape=(2, 3))
+    func = relay.Function([x], relay.reshape(x, (-1,)), attrs=attrs)
+    assert_parses_as(func.astext(), func)
+
+
 if __name__ == "__main__":
     import sys
 

From 95dab5ba7ff92ffc710a9668d13d424cdeba42ba Mon Sep 17 00:00:00 2001
From: Altan Haan <ahaan@octoml.ai>
Date: Tue, 4 May 2021 23:39:39 +0000
Subject: [PATCH 2/2] remove leftover comment

---
 src/parser/parser.cc | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/parser/parser.cc b/src/parser/parser.cc
index d234c329a9f1..793d6bb9a43d 100644
--- a/src/parser/parser.cc
+++ b/src/parser/parser.cc
@@ -1372,7 +1372,6 @@ class Parser {
     Map<String, ObjectRef> kwargs;
     while (Peek()->token_type == TokenType::kIdentifier) {
       auto key = GetHierarchicalName(ParseHierarchicalName().data);
-      // auto key = Match(TokenType::kIdentifier).ToString();
       Match(TokenType::kEqual);
       // TOOD(@jroesch): syntactically what do we allow to appear in attribute right hand side.
       auto value = ParseAttributeValue();