Add list attr recursive (apache#89)

* Add list attr recursive * fix
piiswrong · Dec 29, 2016 · 472ad59 · 472ad59
1 parent f17fea0
commit 472ad59
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/include/nnvm/symbolic.h b/include/nnvm/symbolic.h
@@ -12,6 +12,7 @@
 
 #include <string>
 #include <vector>
+#include <tuple>
 #include <utility>
 
 #include "./base.h"
@@ -168,6 +169,15 @@ class Symbol {
    * \return The created attribute.
    */
   std::unordered_map<std::string, std::string> ListAttrs(ListAttrOption option) const;
+  /*!
+   * \brief Get attribute dictionary from the symbol and all children.
+   *
+   *  For grouped symbol, an error will be raised.
+   *
+   * \return The created attribute in format <operator_name, key, value>.
+   */
+  std::vector<std::tuple<std::string, std::string, std::string> >
+      ListAttrsRecursive() const;
   /*!
    * \brief Create symbolic functor(AtomicSymbol) by given operator and attributes.
    * \param op The operator.

diff --git a/src/core/symbolic.cc b/src/core/symbolic.cc
@@ -472,6 +472,17 @@ std::unordered_map<std::string, std::string> Symbol::ListAttrs(ListAttrOption op
   }
 }
 
+std::vector<std::tuple<std::string, std::string, std::string> >
+    Symbol::ListAttrsRecursive() const {
+  std::vector<std::tuple<std::string, std::string, std::string> > ret;
+  DFSVisit(this->outputs, [&ret](const NodePtr& n) {
+      for (const auto& it : n->attrs.dict) {
+        ret.emplace_back(std::make_tuple(n->attrs.name, it.first, it.second));
+      }
+    });
+  return ret;
+}
+
 Symbol Symbol::CreateFunctor(const Op* op,
                              std::unordered_map<std::string, std::string> attrs) {
   static auto& fnum_vis_output = Op::GetAttr<FNumVisibleOutputs>("FNumVisibleOutputs");