[RELAY][Refactor] TextPrinter, move ret_type after body in Function. (#1918)

include/tvm/attrs.h

@@ -131,6 +131,13 @@ class BaseAttrsNode : public Node {
    */
   inline void PrintDocString(std::ostream &os) const;  // NOLINT(*)
   /*!
+   * \brief Visit attributes that do not equal the default value.
+   *
+   * \note This is useful to extract fields for concise printing.
+   * \param v The visitor
+   */
+  TVM_DLL virtual void VisitNonDefaultAttrs(AttrVisitor* v) = 0;
+  /*!
    * \brief Get the field information
    * \return The fields in the Attrs.
    */
@@ -199,6 +206,7 @@ class DictAttrsNode : public BaseAttrsNode {
   TVM_DLL static Attrs make(Map<std::string, NodeRef> dict);
   // implementations
   void VisitAttrs(AttrVisitor* v) final;
+  void VisitNonDefaultAttrs(AttrVisitor* v) final;
   void InitByPackedArgs(const runtime::TVMArgs& args, bool allow_unknown) final;
   Array<AttrFieldInfo> ListFieldInfo() const final;
   bool ContentEqual(const Node* other) const final;
@@ -300,15 +308,15 @@ struct AttrNopEntry {
     return *this;
   }
   template<typename T>
-  TSelf& set_default(DMLC_ATTRIBUTE_UNUSED T value) {
+  TSelf& set_default(DMLC_ATTRIBUTE_UNUSED const T& value) {
     return *this;
   }
   template<typename T>
-  TSelf& set_lower_bound(DMLC_ATTRIBUTE_UNUSED T begin) {
+  TSelf& set_lower_bound(DMLC_ATTRIBUTE_UNUSED const T& begin) {
     return *this;
   }
   template<typename T>
-  TSelf& set_upper_bound(DMLC_ATTRIBUTE_UNUSED T end) {
+  TSelf& set_upper_bound(DMLC_ATTRIBUTE_UNUSED const T& end) {
     return *this;
   }
 };
@@ -603,7 +611,7 @@ class AttrDocEntry {
     return *this;
   }
   template<typename T>
-  TSelf& set_default(DMLC_ATTRIBUTE_UNUSED T value) {
+  TSelf& set_default(DMLC_ATTRIBUTE_UNUSED const T& value) {
     std::ostringstream os;
     os << info_->type_info << ", default=" << value;
     info_->type_info = os.str();
@@ -649,6 +657,57 @@ class AttrExistVisitor {
     return AttrNopEntry();
   }
 };
+
+template<typename T>
+struct AttrTriggerNonDefaultEntry {
+  using TSelf = AttrTriggerNonDefaultEntry<T>;
+  // constructor
+  AttrTriggerNonDefaultEntry(
+      AttrVisitor* visitor, const char* key, T* data)
+      : visitor_(visitor), key_(key), data_(data) {}
+
+  ~AttrTriggerNonDefaultEntry() DMLC_THROW_EXCEPTION {
+    if (trigger_) {
+      visitor_->Visit(key_, data_);
+    }
+  }
+  TSelf& describe(DMLC_ATTRIBUTE_UNUSED const char* str) {
+    return *this;
+  }
+  TSelf& set_default(const T& value) {
+    if (AttrsEqual()(value, *data_)) {
+      trigger_ = false;
+    }
+    return *this;
+  }
+  TSelf& set_lower_bound(DMLC_ATTRIBUTE_UNUSED const T& begin) {
+    return *this;
+  }
+  TSelf& set_upper_bound(DMLC_ATTRIBUTE_UNUSED const T& end) {
+    return *this;
+  }
+
+ private:
+  AttrVisitor* visitor_;
+  const char * key_;
+  T *data_;
+  bool trigger_{true};
+};
+
+class AttrNonDefaultVisitor {
+ public:
+  explicit AttrNonDefaultVisitor(AttrVisitor* visitor)
+      : visitor_(visitor) {
+  }
+  template<typename T>
+  AttrTriggerNonDefaultEntry<T>
+  operator()(const char* key, T* value) {
+    return AttrTriggerNonDefaultEntry<T>(visitor_, key, value);
+  }
+
+ private:
+  AttrVisitor* visitor_;
+};
 }  // namespace detail
 
 /*!
@@ -665,6 +724,11 @@ class AttrsNode : public BaseAttrsNode {
     self()->__VisitAttrs__(vis);
   }
 
+  void VisitNonDefaultAttrs(AttrVisitor* v) final {
+    detail::AttrNonDefaultVisitor vis(v);
+    self()->__VisitAttrs__(vis);
+  }
+
   void InitByPackedArgs(const runtime::TVMArgs& args, bool allow_unknown) final {
     CHECK_EQ(args.size() % 2, 0);
     const int kLinearSearchBound = 16;

include/tvm/relay/attrs/nn.h

@@ -13,7 +13,7 @@ namespace tvm {
 namespace relay {
 
 /*! \brief Attributes used in convolution operators */
-struct ConvAttrs : public tvm::AttrsNode<ConvAttrs> {
+struct Conv2DAttrs : public tvm::AttrsNode<Conv2DAttrs> {
   Array<IndexExpr> strides;
   Array<IndexExpr> padding;
   Array<IndexExpr> dilation;
@@ -25,7 +25,7 @@ struct ConvAttrs : public tvm::AttrsNode<ConvAttrs> {
   std::string out_layout;
   DataType out_dtype;
 
-  TVM_DECLARE_ATTRS(ConvAttrs, "relay.attrs.ConvAttrs") {
+  TVM_DECLARE_ATTRS(Conv2DAttrs, "relay.attrs.Conv2DAttrs") {
     TVM_ATTR_FIELD(strides).set_default(Array<IndexExpr>({1, 1}))
         .describe("Specifies the strides of the convolution.");
     TVM_ATTR_FIELD(padding).set_default(Array<IndexExpr>({0, 0}))
@@ -55,14 +55,14 @@ struct ConvAttrs : public tvm::AttrsNode<ConvAttrs> {
         .describe("Dimension ordering of weight. Can be 'OIHW', 'OIHW16o16i', etc."
                   "'O', 'I', 'H', 'W' stands for num_filter, input_channel, height, and width"
                   "dimensions respectively.");
-    TVM_ATTR_FIELD(out_layout).set_default("__undef__")
+    TVM_ATTR_FIELD(out_layout).set_default("")
         .describe("Dimension ordering of output. Can be 'NCHW', 'NHWC', etc."
                   "'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
                   "dimensions respectively. Default to be same as input layout.");
 
     // use 0 bits to indicate none.
     TVM_ATTR_FIELD(out_dtype)
-        .set_default(Int(0))
+        .set_default(NullValue<DataType>())
         .describe("Output data type, set to explicit type under mixed precision setting");
   }
 };
@@ -123,7 +123,7 @@ struct Conv2DTransposeAttrs : public tvm::AttrsNode<Conv2DTransposeAttrs> {
                 "'O', 'I', 'H', 'W' stands for num_filter, input_channel, height, and width"
                 "dimensions respectively.");
     TVM_ATTR_FIELD(out_dtype)
-        .set_default(Int(0))
+        .set_default(NullValue<DataType>())
         .describe("Output data type, set to explicit type under mixed precision setting");
   }
 };

include/tvm/relay/attrs/transform.h

@@ -78,7 +78,7 @@ struct InitOpAttrs : public tvm::AttrsNode<InitOpAttrs> {
       .describe("Target shape.");
     TVM_ATTR_FIELD(dtype)
       .describe("Target data type.")
-      .set_default(Int(0));
+      .set_default(NullValue<DataType>());
   }
 };  // struct InitOpAttrs
 

include/tvm/relay/expr.h

@@ -181,8 +181,6 @@ class FunctionNode : public ExprNode {
  public:
   /*! \brief Function parameters */
   tvm::Array<Var> params;
-  /*! \brief User annotated return type of the function. */
-  Type ret_type;
   /*!
    * \brief
    * The expression which represents the computation of the function,
@@ -190,6 +188,8 @@ class FunctionNode : public ExprNode {
    * or sub-expressions may reference the type variables.
    */
   Expr body;
+  /*! \brief User annotated return type of the function. */
+  Type ret_type;
   /*!
    * \brief Type parameters of the function.
    *  Enables the function to vary its type based on these.
@@ -201,8 +201,8 @@ class FunctionNode : public ExprNode {
 
   void VisitAttrs(tvm::AttrVisitor* v) final {
     v->Visit("params", &params);
-    v->Visit("ret_type", &ret_type);
     v->Visit("body", &body);
+    v->Visit("ret_type", &ret_type);
     v->Visit("type_params", &type_params);
     v->Visit("span", &span);
     v->Visit("_checked_type_", &checked_type_);
@@ -217,8 +217,8 @@ class FunctionNode : public ExprNode {
   TVM_DLL FuncType func_type_annotation() const;
 
   TVM_DLL static Function make(tvm::Array<Var> params,
-                               Type ret_type,
                                Expr body,
+                               Type ret_type,
                                tvm::Array<TypeParam> ty_params);
 
   static constexpr const char* _type_key = "relay.Function";

include/tvm/relay/op.h

@@ -48,6 +48,11 @@ class OpNode : public relay::ExprNode {
    */
   std::string attrs_type_key;
   /*!
+   * \brief attribute type index,
+   * this field varies in each run and is not exposed to frontend.
+   */
+  uint32_t attrs_type_index{0};
+  /*!
    * \brief number of input arguments to the operator,
    * -1 means it is variable length
    */
@@ -416,6 +421,7 @@ inline OpRegistry& OpRegistry::set_num_inputs(int32_t n) {  // NOLINT(*)
 inline OpRegistry& OpRegistry::set_attrs_type_key(  // NOLINT(*)
     const std::string& type_key) {
   get()->attrs_type_key = type_key;
+  get()->attrs_type_index = Node::TypeKey2Index(type_key.c_str());
   return *this;
 }
 

python/tvm/_ffi/base.py

@@ -14,13 +14,15 @@ from . import libinfo
 #----------------------------
 if sys.version_info[0] == 3:
     string_types = (str,)
-    numeric_types = (float, int, np.float32, np.int32)
+    integer_types = (int, np.int32)
+    numeric_types = integer_types + (float, np.float32)
     # this function is needed for python3
     # to convert ctypes.char_p .value back to python str
     py_str = lambda x: x.decode('utf-8')
 else:
     string_types = (basestring,)
-    numeric_types = (float, int, long, np.float32, np.int32)
+    integer_types = (int, long, np.int32)
+    numeric_types = integer_types + (float, np.float32)
     py_str = lambda x: x
 
 

python/tvm/relay/__init__.py

-# pylint: disable=wildcard-import, redefined-builtin
+# pylint: disable=wildcard-import, redefined-builtin, invalid-name
 """The Relay IR namespace containing the IR definition and compiler."""
 from . import base
 from . import ty
@@ -19,6 +19,9 @@ from . import image
 # Span
 Span = base.Span
 
+# Env
+Environment = env.Environment
+
 # Type
 Type = ty.Type
 TupleType = ty.TupleType
@@ -40,3 +43,7 @@ Call = expr.Call
 Let = expr.Let
 If = expr.If
 TupleGetItem = expr.TupleGetItem
+
+# helper functions
+var = expr.var
+const = expr.const

python/tvm/relay/base.py

@@ -3,6 +3,7 @@
 from __future__ import absolute_import as _abs
 from .._ffi.node import NodeBase, register_node as _register_tvm_node
 from . import _make
+from . import _expr
 
 NodeBase = NodeBase
 
@@ -20,7 +21,19 @@ def register_relay_node(type_key=None):
     return _register_tvm_node(type_key)
 
 
+class RelayNode(NodeBase):
+    def astext(self):
+        """Get the text format of the expression.
+
+        Returns
+        -------
+        text : str
+            The text format of the expression.
+        """
+        return _expr._text_print(self)
+
+
 @register_relay_node
-class Span(NodeBase):
+class Span(RelayNode):
     def __init__(self, source, lineno, col_offset):
         self.__init_handle_by_constructor__(_make.Span, source, lineno, col_offset)

python/tvm/relay/env.py

 # pylint: disable=no-else-return, unidiomatic-typecheck, undefined-variable, wildcard-import
 """A global environment storing everything needed to interpret or compile a Relay program."""
-from .base import register_relay_node, NodeBase
+from .base import register_relay_node, RelayNode
 from . import _make
 from . import _env
 
-
 @register_relay_node
-class Environment(NodeBase):
+class Environment(RelayNode):
     """The global Relay environment containing functions,
     options and more.
     """

python/tvm/relay/expr.py

 # pylint: disable=no-else-return, unidiomatic-typecheck, invalid-name
 """The expression nodes of Relay."""
 from __future__ import absolute_import
-from .base import NodeBase, register_relay_node
-from . import _expr
+
+import numpy as _np
+from .base import RelayNode, register_relay_node
 from . import _make
+from . import ty as _ty
+from .._ffi import base as _base, node as _node
+from .. import nd as _nd
 from .. import convert
 
 
-class Expr(NodeBase):
+class Expr(RelayNode):
     """The base type for all Relay expressions."""
     @property
     def checked_type(self):
@@ -56,7 +60,7 @@ class Tuple(Expr):
 
 @register_relay_node
 class Var(Expr):
-    """A local variable in Tvm.Relay.
+    """A local variable in Relay.
 
     Local variable can be used to declare input
     arguments to a function, or intermediate variables.
@@ -101,26 +105,26 @@ class Function(Expr):
     params: List[tvm.relay.Var]
         List of input parameters to the function.
 
-    ret_type: tvm.relay.Type
-        The return type annotation of the function.
-
     body: tvm.relay.Expr
         The body of the function.
 
+    ret_type: Optional[tvm.relay.Type]
+        The return type annotation of the function.
+
     type_params: Optional[List[tvm.relay.TypeParam]]
         The additional type parameters, this is only
         used in advanced usecase of template functions.
     """
     def __init__(self,
                  params,
-                 ret_type,
                  body,
+                 ret_type=None,
                  type_params=None):
         if type_params is None:
             type_params = convert([])
 
         self.__init_handle_by_constructor__(
-            _make.Function, params, ret_type, body, type_params)
+            _make.Function, params, body, ret_type, type_params)
 
 
 @register_relay_node
@@ -158,7 +162,7 @@ class Let(Expr):
 
     Parameters
     ----------
-    var: tvm.relay.Var
+    variable: tvm.relay.Var
         The local variable to be bound.
 
     value: tvm.relay.Expr
@@ -167,9 +171,9 @@ class Let(Expr):
     body: tvm.relay.Expr
         The body of the let binding.
     """
-    def __init__(self, var, value, body):
+    def __init__(self, variable, value, body):
         self.__init_handle_by_constructor__(
-            _make.Let, var, value, body)
+            _make.Let, variable, value, body)
 
 
 @register_relay_node
@@ -208,4 +212,105 @@ class TupleGetItem(Expr):
         self.__init_handle_by_constructor__(
             _make.TupleGetItem, tuple_value, index)
 
-debug_print = _expr._debug_print
+
+class TupleWrapper(_node.NodeGeneric):
+    """TupleWrapper.
+
+    This class is a Python wrapper for a Relay tuple of known size.
+    It allows for accessing the fields of the Relay tuple as though
+    it were a Python tuple.
+
+    Parameters
+    ----------
+    tuple_value: tvm.relay.Expr
+        The input tuple
+
+    size: int
+        The size of the tuple.
+    """
+    def __init__(self, tuple_value, size):
+        self.tuple_value = tuple_value
+        self.size = size
+
+    def asnode(self):
+        """Returns the underlying Relay tuple if this wrapper is passed
+        as an argument to an FFI function."""
+
+        return self.tuple_value
+
+    def __getitem__(self, key):
+        return self.tuple_value.fields[key]
+
+    def __len__(self):
+        return len(self.tuple_value.fields)
+
+
+def var(name_hint,
+        type_annotation=None,
+        shape=None,
+        dtype="float32"):
+    """Create a new tvm.relay.Var.
+
+    This is a simple wrapper function that allows specify
+    shape and dtype directly.
+
+    Parameters
+    ----------
+    name_hint: str
+        The name of the variable.
+        This name only acts as a hint, and is not used
+        for equality.
+
+    type_annotation: Optional[tvm.relay.Type, str]
+        The type annotation on the variable.
+        When type_annotation is a str, we will create a scalar variable.
+
+    shape: Optional[List[tvm.Expr]]
+        The shape of the tensor type.
+
+    dtype: str, optional
+        The data type of the tensor.
+
+    Examples
+    --------
+    .. code-block:: python
+
+      # The following 4 lines are equivalent to each other
+      x = tvm.relay.Var("x", tvm.relay.TensorType([1, 2]))
+      x = tvm.relay.var("x", tvm.relay.TensorType([1, 2]))
+      x = tvm.relay.var("x", shape=[1, 2])
+      x = tvm.relay.var("x", shape=[1, 2], dtype="float32")
+
+      # The following 2 lines are equivalent to each other.
+      y = tvm.relay.var("x", "float32")
+      y = tvm.relay.var("x", shape=(), dtype="float32")
+    """
+    if type_annotation is not None and shape is not None:
+        raise ValueError("Can only specify either type_annotation or shape.")
+    if shape is not None:
+        type_annotation = _ty.TensorType(shape, dtype)
+    elif isinstance(type_annotation, str):
+        type_annotation = _ty.TensorType((), type_annotation)
+    return Var(name_hint, type_annotation)
+
+
+def const(value, dtype=None):
+    """Create a constant value.
+
+    Parameters
+    ----------
+    value: Union[bool, int, float, numpy.ndarray, tvm.nd.NDArray]
+        The constant value.
+
+    dtype: str, optional
+        The data type of the value.
+    """
+    if isinstance(value, _base.numeric_types):
+        value = _np.array(value, dtype=dtype)
+    elif isinstance(value, (bool, list)):
+        value = _np.array(value, dtype=dtype)
+    if isinstance(value, (_np.ndarray, _np.generic)):
+        value = _nd.array(value)
+    if not isinstance(value, _nd.NDArray):
+        raise ValueError("value has to be scalar or NDArray")
+    return Constant(value)

python/tvm/relay/ir_builder.py

@@ -11,32 +11,6 @@ from .expr import Expr, Constant, Let, Var, Function, If
 from .env import Environment
 
 
-class TupleWrapper(tvm._ffi.node.NodeGeneric):
-    """TupleWrapper.
-
-    This class is a Python wrapper for a Relay tuple of known size.
-    It allows for accessing the fields of the Relay tuple as though
-    it were a Python tuple.
-    """
-
-    def __init__(self, tuple_value, size):
-        self.tuple_value = tuple_value
-        self.size = size
-
-
-    def asnode(self):
-        """Returns the underlying Relay tuple if this wrapper is passed
-        as an argument to an FFI function."""
-
-        return self.tuple_value
-
-    def __getitem__(self, key):
-        return self.tuple_value.fields[key]
-
-    def __len__(self):
-        return len(self.tuple_value.fields)
-
-
 def _convert_to_value(arg, ctxt=tvm.cpu(0)):
     # type: (Any, tvm.Context) -> tvm.nd.NDArray
     """Convert Python values into the appropriate types
@@ -132,8 +106,8 @@ class PartialFunc(object):
         """Converts a PartialFunc into a :py:class:`~relay.Function`."""
         return Function(
             self.params,
-            self.ret_type,
             self.body,
+            self.ret_type,
             self.type_params)
 
 #pylint: disable=invalid-name
@@ -325,7 +299,7 @@ class IRBuilder(object):
         def _on_exit():
             bindings, _, _, ret_value = self.exit_scope()
             exp = _mk_let(bindings, ret_value)
-            self.env.add(name, Function(params, ret_type, exp))
+            self.env.add(name, Function(params, exp, ret_type))
 
         return WithScope(10, _on_exit)
 

python/tvm/relay/op/nn/nn.py

 """Neural network operations."""
 from __future__ import absolute_import as _abs
-from tvm.relay.ir_builder import TupleWrapper
+from ...expr import TupleWrapper
 from . import _make
 
 
@@ -145,7 +145,7 @@ def conv2d_transpose(data,
                                   weight_layout, output_padding, out_dtype)
 
 
-def softmax(data, axis):
+def softmax(data, axis=1):
     r"""Computes softmax.
 
     .. math:: \text{softmax}(x)_i = \frac{exp(x_i)}{\sum_j exp(x_j)}
@@ -158,7 +158,7 @@ def softmax(data, axis):
     data: relay.Expr
         The input data to the operator.
 
-    axis: int
+    axis: int, optional
         The axis to sum over when computing softmax
 
     Returns

python/tvm/relay/ty.py

 # pylint: disable=no-else-return, unidiomatic-typecheck, invalid-name
 """The type nodes of the Relay language."""
 from enum import IntEnum
-from .base import NodeBase, register_relay_node
+from .base import RelayNode, register_relay_node
 from . import _make
 
 
-class Type(NodeBase):
+class Type(RelayNode):
     """The base type for all Relay types."""
 
     def __eq__(self, other):
@@ -21,27 +21,25 @@ class Type(NodeBase):
         """Compares two Relay types by referential equality."""
         return super().__eq__(other)
 
+
 @register_relay_node
 class TensorType(Type):
-    """A concrete TensorType in Relay, see tvm/relay/type.h for more details.
+    """A concrete TensorType in Relay.
 
     This is the type assigned to tensor's with a known dype and shape. For
     example a tensor of `float32` and `(5, 5)`.
-    """
-
-    def __init__(self, shape, dtype):
-        """Construct a tensor type.
 
-        Parameters
-        ----------
-        shape: list of tvm.Expr
-        dtype: str
+    Parameters
+    ----------
+    shape: List[tvm.Expr]
+        The shape of the Tensor
 
-        Returns
-        -------
-        tensor_type: The TensorType
-        """
-        self.__init_handle_by_constructor__(_make.TensorType, shape, dtype)
+    dtype: str, optional
+        The content data type.
+    """
+    def __init__(self, shape, dtype="float32"):
+        self.__init_handle_by_constructor__(
+            _make.TensorType, shape, dtype)
 
 
 class Kind(IntEnum):

src/lang/attr_functor.h

@@ -17,11 +17,15 @@ namespace tvm {
 template <typename FType>
 class AttrFunctor;
 
+#define ATTR_FUNCTOR_DEFAULT                                        \
+  { return VisitAttrDefault_(op, std::forward<Args>(args)...); }
+
+
 #define ATTR_FUNCTOR_DISPATCH(OP)                                       \
   vtable.template set_dispatch<OP>(                                     \
       [](const NodeRef& n, TSelf* self, Args... args) {                 \
-        return self->Visit_(static_cast<const OP*>(n.node_.get()),      \
-                            std::forward<Args>(args)...);               \
+        return self->VisitAttr_(static_cast<const OP*>(n.node_.get()),  \
+                                std::forward<Args>(args)...);           \
       });                                                               \
 
 // A functor for common attribute information.
@@ -40,21 +44,21 @@ class AttrFunctor<R(const NodeRef& n, Args...)> {
    * \param args Additional arguments.
    * \return The result of the call
    */
-  virtual R Visit(const NodeRef& n, Args... args) {
+  virtual R VisitAttr(const NodeRef& n, Args... args) {
     static FType vtable = InitVTable();
     if (vtable.can_dispatch(n)) {
       return vtable(n, this, std::forward<Args>(args)...);
     } else {
-      return VisitDefault_(n, std::forward<Args>(args)...);
+      return VisitAttrDefault_(n.get(), std::forward<Args>(args)...);
     }
   }
-  virtual R Visit_(const ArrayNode* op, Args... args) = 0;
-  virtual R Visit_(const StrMapNode* op, Args... args) = 0;
-  virtual R Visit_(const ir::IntImm* op, Args... args) = 0;
-  virtual R Visit_(const ir::UIntImm* op, Args... args) = 0;
-  virtual R Visit_(const ir::FloatImm* op, Args... args) = 0;
-  virtual R Visit_(const ir::StringImm* op, Args... args) = 0;
-  virtual R VisitDefault_(const NodeRef& n, Args... args) = 0;
+  virtual R VisitAttr_(const ArrayNode* op, Args... args) ATTR_FUNCTOR_DEFAULT;
+  virtual R VisitAttr_(const StrMapNode* op, Args... args) ATTR_FUNCTOR_DEFAULT;
+  virtual R VisitAttr_(const ir::IntImm* op, Args... args) ATTR_FUNCTOR_DEFAULT;
+  virtual R VisitAttr_(const ir::UIntImm* op, Args... args) ATTR_FUNCTOR_DEFAULT;
+  virtual R VisitAttr_(const ir::FloatImm* op, Args... args) ATTR_FUNCTOR_DEFAULT;
+  virtual R VisitAttr_(const ir::StringImm* op, Args... args) ATTR_FUNCTOR_DEFAULT;
+  virtual R VisitAttrDefault_(const Node* node, Args... args) = 0;
 
  private:
   // initialize the vtable.

src/lang/attrs.cc

@@ -11,6 +11,10 @@ void DictAttrsNode::VisitAttrs(AttrVisitor* v)  {
   v->Visit("__dict__", &dict);
 }
 
+void DictAttrsNode::VisitNonDefaultAttrs(AttrVisitor* v) {
+  v->Visit("__dict__", &dict);
+}
+
 void DictAttrsNode::InitByPackedArgs(
     const runtime::TVMArgs& args, bool allow_unknown) {
   for (int i = 0; i < args.size(); i += 2) {
@@ -55,48 +59,48 @@ class AttrsEqualChecker :
     if (!equal_) return false;
     if (lhs.same_as(rhs)) return true;
     if (!lhs.defined() || !rhs.defined()) return false;
-    if (!this->Visit(lhs, rhs)) {
+    if (!this->VisitAttr(lhs, rhs)) {
       equal_ = false;
     }
     return equal_;
   }
 
-  bool VisitDefault_(const NodeRef& lhs, const NodeRef& other) final {
+  bool VisitAttrDefault_(const Node* lhs, const NodeRef& other) final {
     if (lhs->derived_from<BaseAttrsNode>()) {
-      return static_cast<const BaseAttrsNode*>(lhs.get())->ContentEqual(other.get());
+      return static_cast<const BaseAttrsNode*>(lhs)->ContentEqual(other.get());
     }
-    return lhs.same_as(other);
+    return lhs == other.get();
   }
 
-  bool Visit_(const IntImm* lhs, const NodeRef& other) final {
+  bool VisitAttr_(const IntImm* lhs, const NodeRef& other) final {
     if (const auto* rhs = other.as<IntImm>()) {
       return lhs->value == rhs->value;
     }
     return false;
   }
 
-  bool Visit_(const UIntImm* lhs, const NodeRef& other) final {
+  bool VisitAttr_(const UIntImm* lhs, const NodeRef& other) final {
     if (const auto* rhs = other.as<UIntImm>()) {
       return lhs->value == rhs->value;
     }
     return false;
   }
 
-  bool Visit_(const FloatImm* lhs, const NodeRef& other) final {
+  bool VisitAttr_(const FloatImm* lhs, const NodeRef& other) final {
     if (const auto* rhs = other.as<FloatImm>()) {
       return lhs->value == rhs->value;
     }
     return false;
   }
 
-  bool Visit_(const StringImm* lhs, const NodeRef& other) final {
+  bool VisitAttr_(const StringImm* lhs, const NodeRef& other) final {
     if (const auto* rhs = other.as<StringImm>()) {
       return lhs->value == rhs->value;
     }
     return false;
   }
 
-  bool Visit_(const ArrayNode* lhs, const NodeRef& other) final {
+  bool VisitAttr_(const ArrayNode* lhs, const NodeRef& other) final {
     if (const auto* rhs = other.as<ArrayNode>()) {
       if (rhs->data.size() != lhs->data.size()) return false;
       for (size_t  i = 0; i < lhs->data.size(); ++i) {
@@ -106,7 +110,7 @@ class AttrsEqualChecker :
     return true;
   }
 
-  bool Visit_(const StrMapNode* lhs, const NodeRef& other) final {
+  bool VisitAttr_(const StrMapNode* lhs, const NodeRef& other) final {
     if (const auto* rhs = other.as<StrMapNode>()) {
       if (rhs->data.size() != lhs->data.size()) return false;
       for (const auto& kv : lhs->data) {
@@ -127,38 +131,38 @@ class AttrContentHasher :
  public:
   size_t result_{0};
 
-  void VisitDefault_(const NodeRef& value) final {
+  void VisitAttrDefault_(const Node* value) final {
     if (value->derived_from<BaseAttrsNode>()) {
-      Update(static_cast<const BaseAttrsNode*>(value.get())->ContentHash());
+      Update(static_cast<const BaseAttrsNode*>(value)->ContentHash());
     } else {
-      Update(NodeHash()(value));
+      Update(NodeHash()(GetRef<NodeRef>(value)));
     }
   }
 
-  void Visit_(const IntImm* op) final {
+  void VisitAttr_(const IntImm* op) final {
     Update(std::hash<int64_t>()(op->value));
   }
 
-  void Visit_(const UIntImm* op) final {
+  void VisitAttr_(const UIntImm* op) final {
     Update(std::hash<uint64_t>()(op->value));
   }
 
-  void Visit_(const FloatImm* op) final {
+  void VisitAttr_(const FloatImm* op) final {
     Update(std::hash<double>()(op->value));
   }
 
-  void Visit_(const StringImm* op) final {
+  void VisitAttr_(const StringImm* op) final {
     Update(std::hash<std::string>()(op->value));
   }
 
-  void Visit_(const ArrayNode* op) final {
+  void VisitAttr_(const ArrayNode* op) final {
     Update(op->data.size());
     for (size_t  i = 0; i < op->data.size(); ++i) {
-      this->Visit(NodeRef(op->data[i]));
+      this->VisitAttr(NodeRef(op->data[i]));
     }
   }
 
-  void Visit_(const StrMapNode* lhs) final {
+  void VisitAttr_(const StrMapNode* lhs) final {
     using Entry = std::pair<std::string, NodePtr<Node> >;
     std::vector<Entry> data(lhs->data.begin(), lhs->data.end());
     std::sort(data.begin(), data.end(), [](const Entry& a, const Entry& b) {
@@ -166,7 +170,7 @@ class AttrContentHasher :
       });
     for (const Entry& kv : data) {
       Update(std::hash<std::string>()(kv.first));
-      this->Visit(NodeRef(kv.second));
+      this->VisitAttr(NodeRef(kv.second));
     }
   }
 
@@ -184,7 +188,7 @@ bool AttrsEqual::Equal(const NodeRef& lhs, const NodeRef& rhs) {
 size_t AttrsHash::Hash(const NodeRef& node) {
   if (!node.defined()) return 0;
   AttrContentHasher hasher;
-  hasher.Visit(node);
+  hasher.VisitAttr(node);
   return hasher.result_;
 }
 

src/lang/reflection.cc

@@ -208,6 +208,8 @@ class JSONAttrGetter : public AttrVisitor {
     node_->type_key = node->type_key();
     // sepcially handle global object
     auto* f = dmlc::Registry<NodeFactoryReg>::Find(node_->type_key);
+    CHECK(f != nullptr)
+        << "Node type \'" << node_->type_key << "\' is not registered in TVM";
     if (f->fglobal_key != nullptr) {
       node_->global_key = f->fglobal_key(node);
       return;

src/relay/ir/base.cc

@@ -51,6 +51,8 @@ Span SpanNode::make(SourceName source, int lineno, int col_offset) {
   return Span(n);
 }
 
+TVM_REGISTER_NODE_TYPE(SpanNode);
+
 TVM_REGISTER_API("relay._make.Span")
 .set_body([](TVMArgs args, TVMRetValue *ret) {
     *ret = SpanNode::make(args[0], args[1], args[2]);

src/relay/ir/debug_printer.cc

-/*!
- *  Copyright (c) 2018 by Contributors
- * \file debug_printer.cc
- * \brief A pretty printer for the Relay IR.
- * As we had not determined a formal syntax yet, right now it is only for debug purpose.
- */
-
-#include <tvm/relay/expr_functor.h>
-#include <tvm/relay/environment.h>
-#include <tvm/relay/error.h>
-#include <iostream>
-#include <sstream>
-#include <vector>
-#include <unordered_map>
-#include <string>
-#include <vector>
-#include <iostream>
-#include "../pass/type_functor.h"
-#include "doc.h"
-
-namespace tvm {
-namespace relay {
-
-using namespace tvm::runtime;
-
-Doc KindDocify(TypeParamNode::Kind k) {
-  switch (k) {
-    case TypeParamNode::kShapeVar:
-      return DocOfStr("ShapeVar");
-    case TypeParamNode::kShape:
-      return DocOfStr("Shape");
-    case TypeParamNode::kBaseType:
-      return DocOfStr("BaseType");
-    case TypeParamNode::kType:
-      return DocOfStr("Type");
-    default:
-      LOG(FATAL) << "unreachable code: case not handle in kind";
-      throw;  // log fatal throw but compiler doesnt know
-  }
-}
-
-template<typename T>
-std::vector<Doc> MapDocify(const tvm::Array<T>& arr, const std::function<Doc(const T&)>& f) {
-  std::vector<Doc> vec;
-  for (size_t i = 0; i < arr.size(); ++i) {
-    vec.push_back(f(arr[i]));
-  }
-  return vec;
-}
-
-template<typename T, typename Hash = std::hash<T>, typename Eq = std::equal_to<T>>
-class Counter {
-  std::unordered_map<T, size_t, Hash, Eq> cnt_;
-
- public:
-  Counter() = default;
-  Counter(const Counter&) = delete;
-  size_t operator()(const T& t) {
-    auto v = cnt_.count(t) == 0 ? 0 : cnt_.at(t) + 1;
-    cnt_[t] = v;
-    return v;
-  }
-};
-
-std::string Mangle(const std::string& str, size_t s) {
-  return str + "_" + std::to_string(s);
-  // return s == 0 ? str : str + "_" + std::to_string(s - 1);
-  // the above line look prettier but is dangerous:
-  // suppose we have x, x, x_0. mangling will give x, x_0, x_0!
-  // the save approach give x_0, x_1, x_0_1, and in fact never clash:
-  // stripping _([0-9]*) is invert of mangle under all circumstances.
-  // another problem is we need to prevent Var/TypeParam/GlobalVar clashing each other.
-}
-
-constexpr size_t indent = 2;
-
-struct TypeParamName {
-  bool operator==(const TypeParamName&) const {
-    return true;
-  }
-};
-
-struct mhash {
-  size_t operator()(const ::tvm::relay::TypeParamName&) const noexcept {
-    return 0;
-  }
-};
-
-class TypeDocifier : private TypeFunctor<Doc(const Type& n)> {
-  Environment env;
-  Counter<TypeParamName, mhash> cnt;
-  std::unordered_map<TypeParam, Doc, NodeHash, NodeEqual> map;
-
-  std::vector<Doc> DocifyTypeArray(const tvm::Array<Type>& arr) {
-    return MapDocify<Type>(arr, [=](const Type& t) { return Docify(t); });
-  }
-
-  std::vector<Doc> DocifyTypeParam(const tvm::Array<TypeParam>& arr) {
-    return MapDocify<TypeParam>(arr, [=](const TypeParam& tp) {
-        return Docify(tp);
-      });
-  }
-
-  std::vector<Doc> DocifyTypeConstraint(const tvm::Array<TypeConstraint>& arr) {
-    return MapDocify<TypeConstraint>(arr, [=](const TypeConstraint& tc) { return Docify(tc); });
-  }
-
-  Doc VisitType_(const TensorTypeNode* t) final {
-    return DocOfStr("tensor");
-  }
-
-  Doc VisitType_(const TypeParamNode* p) final {
-    auto tp = GetRef<TypeParam>(p);
-    if (map.count(tp) == 0) {
-      auto name =
-        DocOfStr(Mangle("tp", cnt(TypeParamName())) +
-                 std::string(":")) +
-        KindDocify(p->kind);
-      map.insert(std::pair<TypeParam, Doc>(tp, name));
-    }
-    return map.at(tp);
-  }
-
-  Doc Quantify(const tvm::Array<TypeParam>& tp, const Doc& d) {
-    if (tp.size() == 0) {
-      return d;
-    }
-    return Seq("forall", DocifyTypeParam(tp), ",") + Sep() + d;
-  }
-
-  Doc Constraint(const tvm::Array<TypeConstraint>& tc, const Doc& d) {
-    if (tc.size() == 0) {
-      return d;
-    }
-    return Seq("(", DocifyTypeConstraint(tc), ") =>") + Sep() + d;
-  }
-
-  Doc VisitType_(const FuncTypeNode* f) final {
-    auto inner = Seq("<", DocifyTypeArray(f->arg_types), ">") + Sep() +
-                 DocOfStr("->") + Sep() + Docify(f->ret_type);
-    return Group(Quantify(f->type_params,
-                          Constraint(f->type_constraints, inner)));
-  }
-
-  Doc VisitType_(const TypeRelationNode* r) final {
-    return DocOfStr("Relation") + Seq("(", DocifyTypeArray(r->args), ")");
-  }
-
-  Doc VisitType_(const TupleTypeNode* t) final {
-    return Seq("<", DocifyTypeArray(t->fields), ">");
-  }
-
-  Doc VisitType_(const IncompleteTypeNode* i) final {
-    return DocOfStr("_");
-  }
-
- public:
-  TypeDocifier(const Environment& env) : env(env) { }
-
-  Doc Docify(const Type& t) { return t.get() ? (*this)(t) : DocOfStr("_"); }
-};
-
-class ExprDocifier : private ExprFunctor<Doc(const Expr& n)> {
-  Environment env;
-  Counter<std::string> cnt;
-  std::unordered_map<Var, std::string, NodeHash, NodeEqual> map;
-  TypeDocifier td;
-
-  std::string VarName(const Var& v) {
-    if (map.count(v) == 0) {
-      map.insert(std::pair<Var, std::string>(v, Mangle(v->name_hint, cnt(v->name_hint))));
-    }
-    return map.at(v);
-  }
-
-  Doc TypeAnnotation(const Doc& d, const Type& t) {
-    // test for t being null. probably shouldnt has null. should talk to jared.
-    if (!t.get() || t.as<IncompleteTypeNode>()) {
-      return d;
-    } else {
-      return d + DocOfStr(":") + td.Docify(t);
-    }
-  }
-
-  std::vector<Doc> DocifyExprArray(const tvm::Array<Expr>& arr) {
-    std::vector<Doc> vec;
-    for (size_t i = 0; i < arr.size(); ++i) {
-      vec.push_back(Docify(arr[i]));
-    }
-    return vec;
-  }
-
-  std::vector<Doc> DocifyParamArray(const tvm::Array<Var>& arr) {
-    std::vector<Doc> vec;
-    for (Var param : arr) {
-      vec.emplace_back(TypeAnnotation(DocOfStr(VarName(param)),
-                                      param->type_annotation));
-    }
-    return vec;
-  }
-
-  Doc VisitExpr_(const ConstantNode* c) final {
-    return DocOfStr("some_constant");
-  }
-
-  Doc VisitExpr_(const TupleNode* t) final {
-    return Seq("<", DocifyExprArray(t->fields), ">");
-  }
-
-  Doc VisitExpr_(const VarNode* v) final {
-    return DocOfStr(VarName(GetRef<Var>(v)));
-  }
-
-  Doc VisitExpr_(const GlobalVarNode* g) final {
-    return DocOfStr(g->name_hint);
-  }
-
-  Doc VisitExpr_(const FunctionNode* f) final {
-    return Group(TypeAnnotation(Seq("(", DocifyParamArray(f->params), ")"), f->ret_type) + Sep() +
-                 DocOfStr("=>") + Sep() +
-                 Block(indent, "{", Docify(f->body), "}"));
-  }
-
-  Doc VisitExpr_(const CallNode* c) final {
-    return Docify(c->op) + Seq("<", DocifyExprArray(c->args), ">");
-  }
-
-  Doc VisitExpr_(const LetNode* l) final {
-    return Group(DocOfStr("let") + Sep() +
-                 TypeAnnotation(Docify(l->var), l->var->type_annotation) + Sep() +
-                 DocOfStr("=") + Sep() + Docify(l->value) + DocOfStr(";") + Endl() +
-                 Docify(l->body));
-  }
-
-  Doc VisitExpr_(const IfNode* i) final {
-    return Group(DocOfStr("if") + Sep() + Docify(i->cond) + Sep() +
-                 Block(indent, "{", Docify(i->true_branch), "}") + Sep() +
-                 DocOfStr("else") + Sep() +
-                 Block(indent, "{", Docify(i->false_branch), "}"));
-  }
-
-  Doc VisitExpr_(const OpNode* o) final {
-    return DocOfStr(o->name);
-  }
-
-  Doc VisitExpr_(const TupleGetItemNode* g) final {
-    return Docify(g->tuple) + DocOfStr(std::string(".") + std::to_string(g->index));
-  }
-
- public:
-  ExprDocifier(const Environment& env) : env(env), td(env) { }
-
-  Doc Docify(const Expr& e) { return (*this)(e); }
-};
-
-Doc DocOfExpr(const Environment& env, const Expr& expr) {
-  ExprDocifier d(env);
-  return d.Docify(expr);
-}
-
-Doc DocOfType(const Environment& env, const Type& expr) {
-  TypeDocifier d(env);
-  return d.Docify(expr);
-}
-
-RDoc ExprRDoc(const Environment& env, const Expr& expr) {
-  return Layout(DocOfExpr(env, expr));
-}
-
-RDoc TypeRDoc(const Environment& env, const Type& expr) {
-  return Layout(DocOfType(env, expr));
-}
-
-std::ostream & DebugPrint(const Environment& env, const Expr& e, std::ostream& os) {
-  return os << ExprRDoc(env, e);
-}
-
-std::ostream & DebugPrint(const Environment& env, const Type& t, std::ostream& os) {
-  return os << TypeRDoc(env, t);
-}
-
-std::string PrintExpr(const Environment& env, const Expr& e) {
-  std::stringstream ss;
-  ss << ExprRDoc(env, e);
-  return ss.str();
-}
-
-std::string PrintType(const Environment& env, const Type& t) {
-  std::stringstream ss;
-  ss << TypeRDoc(env, t);
-  return ss.str();
-}
-
-TVM_REGISTER_API("relay._expr._debug_print")
-.set_body([](TVMArgs args, TVMRetValue* ret) {
-    NodeRef x = args[1];
-    if (x.as<TypeNode>()) {
-      *ret = PrintType(args[0], Downcast<Type>(x));
-    } else {
-      *ret = PrintExpr(args[0], Downcast<Expr>(x));
-    }
-  });
-
-}  // namespace relay
-}  // namespace tvm

src/relay/ir/doc.h

-/*!
- *  Copyright (c) 2018 by Contributors
- * \file doc.h
- * \brief A pretty printer DSL for constructing (Doc) and formatting (RDoc) documents.
- *        It is based heavily on Philip Wadler's "A prettier printer."
- *        See https://homepages.inf.ed.ac.uk/wadler/papers/prettier/prettier.pdf
- *        for more details.
- *
- * Since the original paper uses call by value for efficiency, everything doc function is maximally lazy.
- * You can probably yank speed by doing strict analysis and removing some Lazy (if this is bottleneck).
- */
-#ifndef TVM_RELAY_IR_DOC_H_
-#define TVM_RELAY_IR_DOC_H_
-
-#include <tvm/relay/error.h>
-#include <unordered_map>
-#include <utility>
-#include <string>
-#include <functional>
-#include <vector>
-#include <memory>
-#include <ostream>
-#include <map>
-
-namespace tvm {
-namespace relay {
-
-/*! \brief A Document represent structured text.
- * beside having unstructured string, it capture different ways to compose them -
- * line break, space, indentation, representation choice.
- */
-struct Doc;
-
-/*! \brief RDoc represent rendered document.
- * all the high level detail on the document, such as indentation, choice, has been removed.
- * there is only one single, straight forward way to print it.
- */
-struct RDoc;
-
-//! \brief Empty document
-inline Doc Nil();
-
-//! \brief Concatenate two documents
-inline Doc App(const Doc& l, const Doc& r);
-
-//! \brief Indent a document
-inline Doc Nest(size_t width, const Doc& doc);
-
-//! \brief Lift string to a document
-inline Doc DocOfStr(const std::string& text);
-
-//! \brief New line
-inline Doc Endl();
-
-//! \brief Remove all line break from the Document.
-inline Doc Flatten(const Doc& d);
-
-/*! \brief Choose between two possible layouts.
- * assume Flatten(l) == Flatten(r), and l need to be more compact.
- */
-inline Doc Choose(const Doc& l, const Doc& r);
-
-//! \brief Use a single line if possible
-inline Doc Group(const Doc& d);
-
-//! \brief print an RDoc
-inline std::ostream& operator<<(std::ostream& os, const RDoc& rdoc);
-
-/*! \brief Joins a vector of documents with a given separator document
- *  \example Join(["a", "b, "c"], ", ") => "a, b, c"
- *  \param vec the vector of documents
- *  \param sep the separator between documents
- */
-inline Doc Join(const std::vector<Doc>& vec, const Doc& sep);
-
-/*! \brief Creates an indented block.
- *  \param indent the indentation size
- *  \param open the opening string
- *  \param body the body of the block
- *  \param close the closing string
- */
-inline Doc Block(size_t indent, const std::string& open,
-                 const Doc& body, const std::string& close);
-
-/*! \brief Creates a comma-separated sequence with opening and closing strings.
- *  \param open the opening string
- *  \param body the body of the Block
- *  \param close the closing string
- */
-inline Doc Seq(const std::string& open,
-               const std::vector<Doc>& body, const std::string& close);
-
-//! \brief Either a space or a new line
-inline Doc Sep();
-
-/*! \brief Layout a document to a given width
- *  \param d the document to render
- *  \param width the line width
- */
-inline RDoc Layout(const Doc& d, size_t width = 80);
-
-// end of API, start of implementation
-
-template<typename T>
-struct LazyNode {
-  mutable std::function<T()> thunk;
-  explicit LazyNode(const std::function<T()>& thunk) : thunk(thunk) { }
-};
-
-//! \brief denote a value that will be computed (at most once) on need.
-template<typename T>
-struct Lazy {
-  std::shared_ptr<LazyNode<T> > lazy_node;
-  explicit Lazy(const std::function<T()>& thunk) :
-    lazy_node(std::make_shared<LazyNode<T>>(thunk)) { }
-  explicit Lazy(const T& value) : Lazy([=]() { return value; }) { }
-  explicit Lazy(const Lazy<Lazy<T>>& thunk) : Lazy([=]() { return thunk.get().get(); }) { }
-  // calculate the result.
-  // memoize it by replacing the thunk with a constant function which immediate return.
-  T get() const {
-    T res = lazy_node->thunk();
-    lazy_node->thunk = [=]() { return res; };
-    return res;
-  }
-  template<typename R>
-  Lazy<R> map(const std::function<R(const T&)>& func) const {
-    Lazy<T> self(*this);
-    return Lazy<R>([=]() -> R { return func(self.get()); });
-  }
-};
-
-struct NilNode;
-struct AppNode;
-struct NestNode;
-struct TextNode;
-struct LineNode;
-struct ChoiceNode;
-
-/*! \brief The inner representation of Doc.
- * a doc represent structured text,
- * and can be rendered onto screen while keeping the structure.
- */
-struct DocNode {
-  /* a docnode is a union of the below node.
-   * exactly one of them will be non null.
-   * their meaning is denoted by the construction function of the same name.
-   * so for example, the meaning of AppNode is exactly a node construct by App.
-   */
-  std::shared_ptr<NilNode> nil;
-  std::shared_ptr<AppNode> app;
-  std::shared_ptr<NestNode> nest;
-  std::shared_ptr<TextNode> text;  // construct by DocOfStr
-  std::shared_ptr<LineNode> line;
-  std::shared_ptr<ChoiceNode> choice;
-  DocNode(std::shared_ptr<NilNode> nil,
-           std::shared_ptr<AppNode> app,
-           std::shared_ptr<NestNode> nest,
-           std::shared_ptr<TextNode> text,
-           std::shared_ptr<LineNode> line,
-           std::shared_ptr<ChoiceNode> choice) :
-    nil(nil),
-    app(app),
-    nest(nest),
-    text(text),
-    line(line),
-    choice(choice) { }
-};
-
-struct Doc {
-  Lazy<DocNode> doc;
-  explicit Doc(const DocNode& ed) : doc(ed) { }
-  explicit Doc(const Lazy<Doc>& ldoc) :
-    doc(ldoc.map<Lazy<DocNode> >([](const Doc& d){ return d.doc; })) { }
-
-  Doc operator+(const Doc& r) const {
-    return App(*this, r);
-  }
-
-  template<typename T>
-  Lazy<T> Match(
-    const std::function<T()>& nilf,
-    const std::function<T(const Doc&, const Doc&)>& appf,
-    const std::function<T(size_t, const Doc&)>& nestf,
-    const std::function<T(const std::string&)>& textf,
-    const std::function<T()>& linef,
-    const std::function<T(const Doc&, const Doc&)>& choicef) const;
-};
-
-struct NilNode { };
-
-struct AppNode {
-  Doc left, right;
-  AppNode(const Doc& left, const Doc& right) : left(left), right(right) { }
-};
-
-struct NestNode {
-  size_t space;
-  Doc doc;
-  NestNode(size_t space, const Doc& doc) : space(space), doc(doc) { }
-};
-
-struct TextNode {
-  std::string text;
-  explicit TextNode(const std::string& text) : text(text) { }
-};
-
-struct LineNode { };
-
-struct ChoiceNode {
-  Doc left, right;
-  ChoiceNode(const Doc& left, const Doc& right) : left(left), right(right) { }
-};
-
-template<typename T>
-Lazy<T> Doc::Match(
-    const std::function<T()>& nilf,
-    const std::function<T(const Doc&, const Doc&)>& appf,
-    const std::function<T(size_t, const Doc&)>& nestf,
-    const std::function<T(const std::string&)>& textf,
-    const std::function<T()>& linef,
-    const std::function<T(const Doc&, const Doc&)>& choicef) const {
-    return doc.map<T>([=](const DocNode& d) {
-      if (d.nil) {
-        return nilf();
-      } else if (d.app) {
-        return appf(d.app->left, d.app->right);
-      } else if (d.nest) {
-        return nestf(d.nest->space, d.nest->doc);
-      } else if (d.text) {
-        return textf(d.text->text);
-      } else if (d.line) {
-        return linef();
-      } else {
-        return choicef(d.choice->left, d.choice->right);
-      }
-    });
-}
-
-//! \brief Empty document
-inline Doc Nil() {
-  return Doc(DocNode(std::make_shared<NilNode>(), nullptr, nullptr, nullptr, nullptr, nullptr));
-}
-
-//! \brief Concatenate two documents
-inline Doc App(const Doc& l, const Doc& r) {
-  return Doc(DocNode(
-    nullptr,
-    std::make_shared<AppNode>(l, r),
-    nullptr,
-    nullptr,
-    nullptr,
-    nullptr));
-}
-
-//! \brief Indent a document
-inline Doc Nest(size_t width, const Doc& doc) {
-  auto x = std::make_shared<NestNode>(width, doc);
-  return Doc(DocNode(
-    nullptr,
-    nullptr,
-    std::make_shared<NestNode>(width, doc),
-    nullptr,
-    nullptr,
-    nullptr));
-}
-
-//! \brief Lift string to a document
-inline Doc DocOfStr(const std::string& text) {
-  return Doc(DocNode(nullptr, nullptr, nullptr,
-    std::make_shared<TextNode>(text), nullptr, nullptr));
-}
-
-//! \brief New line
-inline Doc Endl() {
-  return Doc(DocNode(nullptr, nullptr, nullptr, nullptr, std::make_shared<LineNode>(), nullptr));
-}
-
-/*! \brief Choose between two possible layouts.
- * assume Flatten(l) == Flatten(r), and l need to be more compact.
- */
-inline Doc Choose(const Doc& l, const Doc& r) {
-  return Doc(DocNode(
-    nullptr,
-    nullptr,
-    nullptr,
-    nullptr,
-    nullptr,
-    std::make_shared<ChoiceNode>(l, r)));
-}
-
-//! \brief Remove new line from the whole document.
-inline Doc Flatten(const Doc& d) {
-  return Doc(d.Match<Doc>(
-    []() { return Nil(); },
-    [](const Doc& l, const Doc& r) { return Flatten(l) + Flatten(r); },
-    [](size_t space, const Doc& doc) { return Flatten(doc); },
-    [](const std::string& str) { return DocOfStr(str); },
-    []() { return DocOfStr(" "); },
-    [](const Doc& l, const Doc& r) { return Flatten(l); }));
-}
-
-//! \brief Use a single line if possible
-inline Doc Group(const Doc& d) {
-  return Choose(Flatten(d), d);
-}
-
-struct RNilNode;
-struct RTextNode;
-struct RLineNode;
-
-struct RDocNode {
-  std::shared_ptr<RNilNode> rnil;
-  std::shared_ptr<RTextNode> rtext;
-  std::shared_ptr<RLineNode> rline;
-  RDocNode(const std::shared_ptr<RNilNode>& rnil,
-           const std::shared_ptr<RTextNode>& rtext,
-           const std::shared_ptr<RLineNode>& rline) :
-    rnil(rnil), rtext(rtext), rline(rline) { }
-};
-
-/*! \brief RDoc represent rendered document.
- * all the high level detail on the document, such as indentation, alternative, has been removed.
- * there is only one single, straight forward way to print it.
- */
-struct RDoc {
-  Lazy<RDocNode> doc;
-  explicit RDoc(const RDocNode& d) : doc(d) { }
-  explicit RDoc(const Lazy<RDoc>& ldoc) :
-    doc(ldoc.map<Lazy<RDocNode>>([](const RDoc& d){ return d.doc; })) { }
-  template<typename T>
-  Lazy<T> Match(
-    const std::function<T()> &rnilf,
-    const std::function<T(const std::string&, const RDoc&)>& rtextf,
-    const std::function<T(size_t, const RDoc&)>& rlinef) const;
-};
-
-inline std::ostream& operator<<(std::ostream& os, const RDoc& rdoc) {
-  return *rdoc.Match<std::ostream*>(
-    [&]() { return & os; },
-    [&](const std::string& text, const RDoc& r) {
-      return & (os << text << r);
-    },
-    [&](size_t space, const RDoc& r) {
-      return & (os << std::endl << std::string(space, ' ') << r);
-    }).get();
-}
-
-struct RNilNode { };
-
-struct RTextNode {
-  std::string text;
-  RDoc rest;
-  RTextNode(const std::string& text, const RDoc& rest) : text(text), rest(rest) { }
-};
-
-struct RLineNode {
-  size_t space;
-  RDoc rest;
-  RLineNode(size_t space, const RDoc& rest) : space(space), rest(rest) { }
-};
-
-//! \brief Empty RDoc
-inline RDoc RNil() { return RDoc(RDocNode(std::make_shared<RNilNode>(), nullptr, nullptr)); }
-
-//! \brief RDoc that begin with std::string
-inline RDoc RText(const std::string& text, const RDoc& rest) {
-  return RDoc(RDocNode(nullptr, std::make_shared<RTextNode>(text, rest), nullptr));
-}
-
-//! \brief RDoc that begin with a new line, followed by space
-inline RDoc RLine(size_t space, const RDoc& rest) {
-  return RDoc(RDocNode(nullptr, nullptr, std::make_shared<RLineNode>(space, rest)));
-}
-
-template<typename T>
-Lazy<T> RDoc::Match(
-  const std::function<T()>& rnilf,
-  const std::function<T(const std::string&, const RDoc&)>& rtextf,
-  const std::function<T(size_t, const RDoc&)>& rlinef) const {
-  return doc.map<T>([=](const RDocNode& rdoc) {
-    if (rdoc.rnil) {
-      return rnilf();
-    } else if (rdoc.rtext) {
-      return rtextf(rdoc.rtext->text, rdoc.rtext->rest);
-    } else {
-      return rlinef(rdoc.rline->space, rdoc.rline->rest);
-    }
-  });
-}
-
-template<typename T>
-struct List;
-
-template<typename T>
-struct EagerList {
-  const std::shared_ptr<std::pair<T, List<T>>> cons;
-};
-
-//! \brief lazy list
-template<typename T>
-struct List {
-  Lazy<EagerList<T> > l;
-  List() : l([]() { return EagerList<T>({nullptr}); }) { }
-  List(const T& t, const List<T>& l) :
-    l([=]() { return EagerList<T>({std::make_shared<std::pair<T, List<T>>>(t, l)}); }) { }
-  template<typename R>
-  Lazy<R> Match(const std::function<R()>& nullf,
-                const std::function<R(const T&, const List<T>&)>& consf) const {
-    return l.template map<R>([=](const EagerList<T>& l) {
-        if (l.cons) {
-          return consf(l.cons->first, l.cons->second);
-        } else {
-          return nullf();
-        }
-    });
-  }
-};
-
-//! \brief Does x fit into line of size w?
-inline bool Fits(int w, const RDoc& x) {
-  return (w >= 0) && x.Match<bool>(
-    []() { return true; },
-    [=](const std::string& s, const RDoc& x) { return Fits(w - s.size(), x); },
-    [](size_t space, const RDoc& x) { return true; }).get();
-}
-
-//! \brief Choose the one that fits best.
-inline RDoc Better(size_t w, size_t k, const RDoc& x, const RDoc& y) {
-  return Fits(w-k, x) ? x : y;
-}
-
-typedef std::pair<size_t/*indent size*/, Doc> best_arg;
-inline RDoc Best(size_t w/*wrap width*/, size_t k/*space used*/,
-  const List<best_arg>& l/*to be rendered*/) {
-  return RDoc(l.Match<RDoc>(
-    []() { return RNil(); },
-    [=](const best_arg& p, const List<best_arg>& z) {
-      return RDoc(p.second.Match<RDoc>(
-        [=]() { return Best(w, k, z); },
-        [=](const Doc& x, const Doc& y) {
-          return Best(
-            w,
-            k,
-            List<best_arg>(best_arg(p.first, x), List<best_arg>(best_arg(p.first, y), z))); },
-        [=](size_t j, const Doc& x) {
-          return Best(w, k, List<best_arg>(best_arg(p.first + j, x), z)); },
-        [=](const std::string& text) { return RText(text, Best(w, k + text.size(), z)); },
-        [=]() { return RLine(p.first, Best(w, p.first, z)); },
-        [=](const Doc& x, const Doc& y) {
-          return Better(
-            w,
-            k,
-            Best(w, k, List<best_arg>(best_arg(p.first, x), z)),
-            Best(w, k, List<best_arg>(best_arg(p.first, y), z))); }));
-    }));
-}
-
-/*! \brief Joins a vector of documents with a given separator document
- *  \example Join(["a", "b, "c"], ", ") => "a, b, c"
- *  \param vec the vector of documents
- *  \param sep the separator between documents
- */
-inline Doc Join(const std::vector<Doc>& vec, const Doc& sep) {
-  // https://www.safaribooksonline.com/library/view/c-cookbook/0596007612/ch04s09.html
-  Doc output = Nil();
-  for (auto p = vec.begin(); p != vec.end(); ++p) {
-    output = output + *p;
-    if (p != vec.end() - 1) {
-      output = output + sep;
-    }
-  }
-
-  return output;
-}
-
-/*! \brief Creates an indented block.
- *  \param indent the indentation size
- *  \param open the opening string
- *  \param body the body of the block
- *  \param close the closing string
- */
-inline Doc Block(size_t indent, const std::string& open,
-  const Doc& body, const std::string& close) {
-  return DocOfStr(open) + Nest(indent, Endl() + body) + Endl() + DocOfStr(close);
-}
-
-/*! \brief Creates a comma-separated sequence with opening and closing strings.
- *  \param open the opening string
- *  \param body the body of the Block
- *  \param close the closing string
- */
-inline Doc Seq(const std::string& open,
-  const std::vector<Doc>& body, const std::string& close) {
-  return Group(DocOfStr(open) +
-               Nest(open.size(), Join(body, DocOfStr(",") + Endl())) +
-               DocOfStr(close));
-}
-
-//! \brief Either a space or a new line
-inline Doc Sep() {
-  return Choose(DocOfStr(" "), Endl());
-}
-
-/*! \brief Layout a document to a given width
- *  \param d the document to render
- *  \param width the line width
- */
-inline RDoc Layout(const Doc& d, size_t width) {
-  return Best(width, 0, List<best_arg>(best_arg(0, d), List<best_arg>()));
-}
-
-}  // namespace relay
-}  // namespace tvm
-#endif  // TVM_RELAY_IR_DOC_H_

src/relay/ir/environment.cc

@@ -100,6 +100,8 @@ void EnvironmentNode::Merge(const Environment &env) {
   }
 }
 
+TVM_REGISTER_NODE_TYPE(EnvironmentNode);
+
 TVM_REGISTER_API("relay._make.Environment")
 .set_body([](TVMArgs args, TVMRetValue *ret) {
     *ret = EnvironmentNode::make(args[0]);

src/relay/ir/expr.cc

@@ -17,6 +17,8 @@ Constant ConstantNode::make(runtime::NDArray data) {
   return Constant(n);
 }
 
+TVM_REGISTER_NODE_TYPE(ConstantNode);
+
 TVM_REGISTER_API("relay._make.Constant")
 .set_body([](TVMArgs args, TVMRetValue *ret) {
     *ret = ConstantNode::make(args[0]);
@@ -44,6 +46,8 @@ Tuple TupleNode::make(tvm::Array<relay::Expr> fields) {
   return Tuple(n);
 }
 
+TVM_REGISTER_NODE_TYPE(TupleNode);
+
 TVM_REGISTER_API("relay._make.Tuple")
 .set_body([](TVMArgs args, TVMRetValue *ret) {
     *ret = TupleNode::make(args[0]);
@@ -61,6 +65,8 @@ Var VarNode::make(std::string name_hint, Type type_annotation) {
   return Var(n);
 }
 
+TVM_REGISTER_NODE_TYPE(VarNode);
+
 TVM_REGISTER_API("relay._make.Var")
 .set_body([](TVMArgs args, TVMRetValue *ret) {
     *ret = VarNode::make(args[0], args[1]);
@@ -82,6 +88,8 @@ GlobalVar GlobalVarNode::make(std::string name_hint) {
   return GlobalVar(n);
 }
 
+TVM_REGISTER_NODE_TYPE(GlobalVarNode);
+
 TVM_REGISTER_API("relay._make.GlobalVar")
 .set_body([](TVMArgs args, TVMRetValue *ret) {
     *ret = GlobalVarNode::make(args[0]);
@@ -94,13 +102,13 @@ TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable)
 
 
 Function FunctionNode::make(tvm::Array<Var> params,
-                            Type ret_type,
                             Expr body,
+                            Type ret_type,
                             tvm::Array<TypeParam> type_params) {
   NodePtr<FunctionNode> n = make_node<FunctionNode>();
   n->params = std::move(params);
-  n->ret_type = std::move(ret_type);
   n->body = std::move(body);
+  n->ret_type = std::move(ret_type);
   n->type_params = std::move(type_params);
   return Function(n);
 }
@@ -113,6 +121,8 @@ FuncType FunctionNode::func_type_annotation() const {
   return FuncTypeNode::make(param_types, this->ret_type, this->type_params, {});
 }
 
+TVM_REGISTER_NODE_TYPE(FunctionNode);
+
 TVM_REGISTER_API("relay._make.Function")
 .set_body([](TVMArgs args, TVMRetValue *ret) {
   *ret = FunctionNode::make(args[0], args[1], args[2], args[3]);
@@ -135,6 +145,8 @@ Call CallNode::make(Expr op, Array<Expr> args, Attrs attrs,
   return Call(n);
 }
 
+TVM_REGISTER_NODE_TYPE(CallNode);
+
 TVM_REGISTER_API("relay._make.Call")
 .set_body([](TVMArgs args, TVMRetValue *ret) {
   *ret = CallNode::make(args[0], args[1], args[2], args[3]);
@@ -154,6 +166,8 @@ Let LetNode::make(Var var, Expr value, Expr body) {
   return Let(n);
 }
 
+TVM_REGISTER_NODE_TYPE(LetNode);
+
 TVM_REGISTER_API("relay._make.Let")
 .set_body([](TVMArgs args, TVMRetValue *ret) {
     *ret = LetNode::make(args[0], args[1], args[2]);
@@ -173,6 +187,8 @@ If IfNode::make(Expr cond, Expr true_branch, Expr false_branch) {
   return If(n);
 }
 
+TVM_REGISTER_NODE_TYPE(IfNode);
+
 TVM_REGISTER_API("relay._make.If").set_body([](TVMArgs args, TVMRetValue *ret) {
   *ret = IfNode::make(args[0], args[1], args[2]);
 });
@@ -190,6 +206,8 @@ TupleGetItem TupleGetItemNode::make(Expr tuple, int index) {
   return TupleGetItem(n);
 }
 
+TVM_REGISTER_NODE_TYPE(TupleGetItemNode);
+
 TVM_REGISTER_API("relay._make.TupleGetItem").set_body([](TVMArgs args, TVMRetValue* ret) {
   *ret = TupleGetItemNode::make(args[0], args[1]);
 });

src/relay/ir/expr_functor.cc

@@ -92,7 +92,7 @@ Expr ExprMutator::VisitExpr_(const FunctionNode* op) {
       body.same_as(op->body)) {
     return GetRef<Expr>(op);
   } else {
-    return FunctionNode::make(params, ret_type, body, ty_params);
+    return FunctionNode::make(params, body, ret_type, ty_params);
   }
 }
 

src/relay/ir/text_printer.cc

+/*!
+ *  Copyright (c) 2018 by Contributors
+ * \file text_printer.cc
+ * \brief Text printer to print relay in text form.
+ */
+#include <tvm/relay/environment.h>
+#include <tvm/relay/expr_functor.h>
+#include <sstream>
+#include "../pass/type_functor.h"
+#include "../../lang/attr_functor.h"
+
+namespace tvm {
+namespace relay {
+
+/*!
+ * \brief the text value used in text printer.
+ * Defined as a struct for future compatibility reason
+ */
+struct TextValue {
+  /*! \brief The str representation */
+  std::string name;
+  // constructor
+  TextValue() {}
+  // constructor
+  explicit TextValue(std::string name) : name(name) {}
+};
+
+// operator overloading
+inline std::ostream& operator<<(std::ostream& os, const TextValue& val) {  // NOLINT(*)
+  return os << val.name;
+}
+
+/*!
+ * \brief Meta data context for TextPrinter.
+ *
+ * This is an important part to enable bi-directional serializability.
+ * We use tvm's Node system to build the current IR.
+ * It can be hard to design a text format for all the possible nodes
+ * as the set of nodes can grow when we do more extensions.
+ *
+ * Instead of trying to design readable text format for every nodes,
+ * we support a meta-data section in the text format.
+ * We allow the text format to refer to a node in the meta-data section.
+ *
+ * The meta-data section is a json serialized string of an Array<NodeRef>.
+ * Each element in the meta-data section can be referenced by the text format.
+ * Each meta data node is printed in the following format.
+ *
+ * meta.<type-key-of-node>(<index-in-meta-section>)
+ *
+ * Specifically, consider the following IR(constructed by python).
+ *
+ * \code
+ *
+ * n = tvm.var("n")
+ * x = tvm.relay.var("x", shape=(n, 1))
+ * f = tvm.relay.Function([x], x)
+ * print(f.astext())
+ *
+ * \endcode
+ *
+ * The corresponding text format is shown in the following code block.
+ *
+ * \code
+ *
+ * function(%x: Tensor[(meta.Variable(id=0),), float32]) {
+ *   %x
+ * }
+ * # Meta data section is a json-serialized string
+ * # of the following array.
+ * # [tvm.var("n")]
+ *
+ * \endcode
+ *
+ * Note that we store tvm.var("n") in the meta data section.
+ * Since it is stored in the index-0 in the meta-data seciton,
+ * we print it as meta.Variable(0).
+ *
+ * The text parser can recover this object by loading from the corresponding
+ * location in the meta data section.
+ *
+ * This is is a design trade-off.
+ * It allows us to embedded any meta-data in the text format,
+ * while still being able to tweak the text part of the printed IR easily.
+ */
+class TextMetaDataContext {
+ public:
+  /*!
+   * \brief Get text representation of meta node.
+   * \param node The node to be converted to meta node.
+   * \return A string representation of the meta node.
+   */
+  std::string GetMetaNode(const NodeRef& node) {
+    std::ostringstream os;
+    auto it = meta_index_.find(node);
+    int64_t index;
+    if (it != meta_index_.end()) {
+      index = it->second;
+    } else {
+      index = static_cast<int64_t>(meta_data_.size());
+      meta_data_.push_back(node);
+      meta_index_[node] = index;
+    }
+    os << "meta." << node->type_key() << "(id=" << index << ")";
+    return os.str();
+  }
+  /*!
+   * \brief Get the metadata section in json format.
+   * \return the meta datastring.
+   */
+  std::string GetMetaSection() const {
+    if (meta_data_.size() == 0) return std::string();
+    return SaveJSON(Array<NodeRef>(meta_data_));
+  }
+
+ private:
+  /*! \brief additional metadata stored in TVM json format */
+  std::vector<NodeRef> meta_data_;
+  /*! \brief map from meta data into its index */
+  std::unordered_map<NodeRef, int64_t, NodeHash, NodeEqual> meta_index_;
+};
+
+class TextPrinter :
+    public ExprFunctor<TextValue(const Expr&)> ,
+    public TypeFunctor<void (const Type&, std::ostream& os)>,  // NOLINT(*)
+    public AttrFunctor<void (const NodeRef&, std::ostream& os)> { // NOLINT(*)
+ public:
+  /*!
+   * \brief Print a node to string.
+   * \param node.
+   * \return The string representation.
+   */
+  std::string Print(const NodeRef& node) {
+    if (node.as<FunctionNode>()) {
+      this->PrintFunc(Downcast<Function>(node));
+    } else if (node.as<EnvironmentNode>()) {
+      this->PrintEnv(Downcast<Environment>(node));
+    } else if (node.as_derived<TypeNode>()) {
+      this->PrintType(Downcast<Type>(node), stream_);
+    } else if (node.as_derived<ExprNode>()) {
+      this->PrintExpr(Downcast<Expr>(node));
+    } else {
+      stream_ << node;
+    }
+    std::string meta_json = meta_.GetMetaSection();
+    if (meta_json.length() != 0) {
+      // append meta data in the end.
+      stream_ << "# meta data\n"
+              << "r\"\"\"\n"
+              << meta_json << "\n"
+              << "\"\"\"";
+    }
+    return stream_.str();
+  }
+
+  void PrintFunc(const Function& func) {
+    this->PrintFuncInternal("function", func);
+    stream_ << "\n";
+  }
+
+  void PrintEnv(const Environment& env) {
+    int counter = 0;
+    for (const auto& kv : env->functions) {
+      std::ostringstream os;
+      if (counter++ != 0) {
+        stream_ << "\n";
+      }
+      os << "def @" << kv.first->name_hint;
+      this->PrintFuncInternal(os.str(), kv.second);
+      stream_ << "\n";
+    }
+  }
+
+  void PrintExpr(const Expr& expr) {
+    TextValue val = GetValue(expr);
+    stream_ << val << "\n";
+  }
+
+  /*!
+   * \brief Get text representation of expr.
+   *
+   * This function may generate additional instructions
+   * in order to compute the final result id of expr.
+   *
+   * When trying to recursively print out an Expr.
+   * The caller should always call GetValue of its children first.
+   * Then the caller can print out to stream_ using the obtained value.
+   *
+   * This is to avoid the call of subsequent GetValue print out
+   * additional instructions which get mixed with the partial instruction
+   * printed by the caller.
+   *
+   * \param expr The input expression.
+   * \return The text value of Expr.
+   */
+  TextValue GetValue(const Expr& expr) {
+    auto it = memo_.find(expr);
+    if (it != memo_.end()) return it->second;
+    TextValue val = this->VisitExpr(expr);
+    memo_[expr] = val;
+    return val;
+  }
+  //------------------------------------
+  // Overload of Expr printing functions
+  //------------------------------------
+  TextValue VisitExpr_(const ConstantNode* op) final {
+    // Print out simple scalar directly.
+    if (op->is_scalar()) {
+      std::ostringstream os;
+      DataType dtype = TVMType2Type(op->data->dtype);
+      CHECK_EQ(op->data->ctx.device_type, kDLCPU);
+      if (dtype == Int(32)) {
+        return ConstScalar(dtype, static_cast<const int32_t*>(op->data->data));
+      } else if (dtype == Int(64)) {
+        return ConstScalar(dtype, static_cast<const int64_t*>(op->data->data));
+      } else if (dtype == Float(32)) {
+        return ConstScalar(dtype, static_cast<const float*>(op->data->data));
+      } else if (dtype == Float(64)) {
+        return ConstScalar(dtype, static_cast<const double*>(op->data->data));
+      }
+    }
+    // default fall-back, record it as meta node.
+    TextValue id = this->AllocTempVar();
+    this->PrintIndent();
+    stream_ << id << " = " << meta_.GetMetaNode(GetRef<NodeRef>(op));
+    this->PrintEndInst("\n");
+    return id;
+  }
+
+  TextValue VisitExpr_(const TupleNode* op) final {
+    std::vector<TextValue> fields;
+    for (Expr field : op->fields) {
+      fields.push_back(GetValue(field));
+    }
+    // NOTE: always recursively visit to get ids,
+    // before print out the current line
+    TextValue id = this->AllocTempVar();
+    this->PrintIndent();
+    stream_ << id << " = (";
+    for (size_t i = 0; i < fields.size(); ++i) {
+      stream_ << fields[i];
+      if (i + 1 != fields.size()) {
+        stream_ << ", ";
+      }
+    }
+    stream_ << ')';
+    this->PrintEndInst("\n");
+    return id;
+  }
+
+  TextValue VisitExpr_(const VarNode* op) final {
+    Var var = GetRef<Var>(op);
+    // This is an unbounded var.
+    TextValue val = AllocVarName(var);
+    this->PrintIndent();
+    stream_ << "free_var ";
+    this->PrintVarDecl(var, stream_);
+    this->PrintEndInst("\n");
+    return val;
+  }
+
+  TextValue VisitExpr_(const GlobalVarNode* op) final {
+    return TextValue('@' + op->name_hint);
+  }
+
+  TextValue VisitExpr_(const FunctionNode* op) final {
+    TextValue id = AllocTempVar();
+    std::ostringstream os;
+    os << id << " = function";
+    this->PrintFuncInternal(os.str(), GetRef<Function>(op));
+    this->PrintEndInst("\n");
+    return id;
+  }
+
+  TextValue VisitExpr_(const CallNode* op) final {
+    // TODO(tqchen, M.K.): support generic call
+    // possibly through meta-data
+    CHECK_EQ(op->type_args.size(), 0U)
+        << "generic call not yet supported";
+    TextValue call_op = GetValue(op->op);
+    std::vector<TextValue> args;
+    for (Expr arg : op->args) {
+      args.emplace_back(GetValue(arg));
+    }
+    TextValue id = this->AllocTempVar();
+    this->PrintIndent();
+    stream_ << id << " = " << call_op << "(";
+    for (size_t i = 0; i < args.size(); ++i) {
+      stream_ << args[i];
+      if (i + 1 != args.size()) {
+        stream_ << ", ";
+      }
+    }
+    this->PrintCallAttrs(op->op, op->attrs, stream_);
+    stream_ << ")";
+    this->PrintEndInst("");
+    this->PrintOptionalInfo(GetRef<Expr>(op));
+    stream_ << '\n';
+    return id;
+  }
+
+  TextValue VisitExpr_(const LetNode* op) final {
+    TextValue id = this->AllocTempVar();
+    this->PrintIndent();
+    stream_ << id << " = ";
+    this->PrintScope(GetRef<Expr>(op));
+    this->PrintEndInst("\n");
+    return id;
+  }
+
+  TextValue VisitExpr_(const IfNode* op) final {
+    TextValue id = this->AllocTempVar();
+    this->PrintIndent();
+    stream_ << id << " = ";
+    this->PrintScope(GetRef<Expr>(op));
+    this->PrintEndInst("\n");
+    return id;
+  }
+
+  TextValue VisitExpr_(const OpNode* op) final {
+    return TextValue(op->name);
+  }
+
+  TextValue VisitExpr_(const TupleGetItemNode* op) final {
+    TextValue tuple = GetValue(op->tuple);
+    TextValue id = this->AllocTempVar();
+    this->PrintIndent();
+    stream_ << id << " = " << tuple << "[" << op->index << "]";
+    this->PrintEndInst("\n");
+    return id;
+  }
+
+  /*!
+   * \brief Print the type to os
+   * \param type The type to be printed.
+   * \param os The output type.
+   */
+  void PrintType(const Type& type, std::ostream& os) {  // NOLINT(*)
+    this->VisitType(type, os);
+  }
+  //------------------------------------
+  // Overload of Expr printing functions
+  //------------------------------------
+  void VisitType_(const TensorTypeNode* node, std::ostream& os) final {  // NOLINT(*)
+    // scalar type
+    if (node->shape.size() == 0) {
+      os << runtime::TVMType2String(Type2TVMType(node->dtype));
+      return;
+    }
+    os << "Tensor[(";
+    for (size_t i = 0; i < node->shape.size(); ++i) {
+      this->PrintAttr(node->shape[i], os);
+      if (i + 1 != node->shape.size()) {
+        os << ", ";
+      }
+    }
+    // conform to python tuple format (1,)
+    if (node->shape.size() == 1) {
+      os << ",";
+    }
+    os << "), " << runtime::TVMType2String(Type2TVMType(node->dtype)) << "]";
+  }
+
+  void VisitTypeDefault_(const Node* node, std::ostream& os) final {  // NOLINT(*)
+    // by default always print as meta-data
+    os << meta_.GetMetaNode(GetRef<NodeRef>(node));
+  }
+
+  /*!
+   * \brief Print an attribute value to os.
+   * \param value The value to be printed.
+   * \param os The output type.
+   */
+  void PrintAttr(const NodeRef& value, std::ostream& os) {  // NOLINT(*)
+    this->VisitAttr(value, os);
+  }
+  //------------------------------------
+  // Overload of Attr printing functions
+  //------------------------------------
+  void VisitAttr_(const ArrayNode* op, std::ostream& os) final {  // NOLINT(*)
+    os << "[";
+    for (size_t i = 0; i < op->data.size(); ++i) {
+      this->PrintAttr(NodeRef(op->data[i]), os);
+      if (i + 1 != op->data.size()) {
+        os << ", ";
+      }
+    }
+    os << "]";
+  }
+  void VisitAttrDefault_(const Node* op, std::ostream& os) final { // NOLINT(*)
+    os << meta_.GetMetaNode(GetRef<NodeRef>(op));
+  }
+
+  void VisitAttr_(const ir::IntImm* op, std::ostream& os) final {  // NOLINT(*)
+    this->PrintConstScalar(op->type, &(op->value), os);
+  }
+
+  void VisitAttr_(const ir::UIntImm* op, std::ostream& os) final {  // NOLINT(*)
+    this->PrintConstScalar(op->type, &(op->value), os);
+  }
+
+  void VisitAttr_(const ir::FloatImm* op, std::ostream& os) final {  // NOLINT(*)
+    this->PrintConstScalar(op->type, &(op->value), os);
+  }
+
+  void VisitAttr_(const ir::StringImm* op, std::ostream& os) final {  // NOLINT(*)
+    this->PrintString(op->value, os);
+  }
+
+ protected:
+  /*!
+   * \brief Print attributes after call.
+   * \param op The operator to be called.
+   * \param attrs The attributes.
+   * \param os The output stream.
+   */
+  void PrintCallAttrs(const Expr& op, const Attrs& attrs, std::ostream& os);  // NOLINT(*)
+
+  /*!
+   * \brief Print the a new scopr.
+   * \param body The body.
+   */
+  void PrintScope(Expr body) {
+    stream_ << "{\n";
+    int sid = this->BeginScope();
+    this->PrintScopeBody(body);
+    this->EndScope(sid);
+    this->PrintIndent();
+    stream_ << "}";
+  }
+  /*!
+   * \brief Print the body of a new scope without {}
+   *
+   * This function will keep printing continuous sequence
+   * of let/if scope without introducing a new scope in the text.
+   *
+   * \param body The body.
+   */
+  void PrintScopeBody(Expr body) {
+    if (const LetNode* let = body.as<LetNode>()) {
+      TextValue value = GetValue(let->value);
+      AllocVarName(let->var);
+      // let var = value;
+      this->PrintIndent();
+      stream_ << "let ";
+      this->PrintVarDecl(let->var, stream_);
+      stream_ << " = " << value;
+      this->PrintEndInst("\n");
+      this->PrintScopeBody(let->body);
+    } else if (const IfNode* ifnode = body.as<IfNode>()) {
+      TextValue cond = GetValue(ifnode->cond);
+      this->PrintIndent();
+      stream_ << "if (" << cond << ") ";
+      this->PrintScope(ifnode->true_branch);
+      this->PrintIndent();
+      stream_ << "else ";
+      this->PrintScope(ifnode->false_branch);
+      this->PrintEndInst("\n");
+    } else {
+      TextValue value = GetValue(body);
+      this->PrintIndent();
+      stream_ << value;
+      this->PrintEndInst("\n");
+    }
+  }
+
+  /*!
+   * \brief Internal function to print a function argument list and its body.
+   * \param prefix The prefix before argument list.
+   * \param fn The function to be printed.
+   */
+  void PrintFuncInternal(std::string prefix, const Function& fn) {
+    // TODO(tqchen, M.K.) support generic function
+    // Possibly through meta-data
+    CHECK_EQ(fn->type_params.size(), 0U)
+        << "generic fn not yet supported";
+    this->PrintIndent();
+    stream_ << prefix << "(";
+    size_t decl_indent = prefix.length() + 1;
+    for (size_t i = 0; i < fn->params.size(); ++i) {
+      if (i != 0) {
+        this->PrintIndent(decl_indent);
+      }
+      AllocVarName(fn->params[i]);
+      this->PrintVarDecl(fn->params[i], stream_);
+      if (i + 1 != fn->params.size()) {
+        stream_ << ",\n";
+      }
+    }
+    stream_ << ") ";
+    if (fn->ret_type.defined()) {
+      stream_ << " -> ";
+      this->PrintType(fn->ret_type, stream_);
+    }
+    this->PrintScope(fn->body);
+  }
+  /*!
+   * \brief Print additional info about expr in comment.
+   * \param expr The expression.
+   */
+  void PrintOptionalInfo(const Expr& expr) {
+    // additional information in comment.
+    if (expr->checked_type_.defined()) {
+      stream_ << " # ty=";
+      this->PrintType(expr->checked_type(), stream_);
+    }
+  }
+  /*!
+   * \brief print var_name[:type]
+   * \param var The variable to be printed
+   * \param os The output stream
+   */
+  void PrintVarDecl(const Var& var, std::ostream& os) {  // NOLINT(*)
+    TextValue v = GetValue(var);
+    os << v;
+    if (var->type_annotation.defined()) {
+      os << ": ";
+      this->PrintType(var->type_annotation, os);
+    }
+  }
+  /*!
+   * \brief Get a constant scalar value.
+   * \param dtype The data type.
+   * \param data The pointer to the data.
+   * \tparam T the content data type holding the data.
+   */
+  template<typename T>
+  TextValue ConstScalar(DataType dtype, const T* data) {
+    std::ostringstream os;
+    PrintConstScalar(dtype, data, os);
+    return TextValue(os.str());
+  }
+  /*!
+   * \brief special method to print out const scalar
+   * \param dtype The data type
+   * \param data The pointer to hold the data.
+   * \param os The output stream.
+   */
+  template<typename T>
+  void PrintConstScalar(DataType dtype, const T* data, std::ostream& os) {  // NOLINT(*)
+    if (dtype == Int(32)) {
+      os << data[0];
+    } else if (dtype == Float(32)) {
+      os << data[0] << 'f';
+    } else if (dtype == Bool()) {
+      PrintBool(data[0] != 0, os);
+    } else {
+      os << dtype << "(" << data[0] << ")";
+    }
+  }
+  /*!
+   * \brief Print constant bool value.
+   * \param value The value to be printed.
+   * \param os The output stream
+   */
+  void PrintBool(bool value, std::ostream& os) { // NOLINT(*)
+    if (value) {
+      os << "True";
+    } else {
+      os << "False";
+    }
+  }
+  /*!
+   * \brief Print constant string.
+   * \param value The value to be printed.
+   * \param os The output stream
+   */
+  void PrintString(const std::string& value, std::ostream& os) { // NOLINT(*)
+    // TODO(M.K.): add escape.
+    os << "\"" << value << "\"";
+  }
+  /*!
+   * \brief get a unique name with the corresponding prefix
+   * \param prefix The prefix of the name
+   * \return The returned name.
+   */
+  std::string GetUniqueName(std::string prefix) {
+    auto it = name_alloc_map_.find(prefix);
+    if (it != name_alloc_map_.end()) {
+      while (true) {
+        std::ostringstream os;
+        os << prefix << (++it->second);
+        std::string name = os.str();
+        if (name_alloc_map_.count(name) == 0) {
+          prefix = name;
+          break;
+        }
+      }
+    }
+    name_alloc_map_[prefix] = 0;
+    return prefix;
+  }
+  /*!
+   * \brief mark the beginning of a new scope
+   * \return The scope id.
+   */
+  int BeginScope() {
+    int sid = static_cast<int>(scope_valid_.size());
+    scope_valid_.push_back(true);
+    indent_ += 2;
+    return sid;
+  }
+  /*!
+   * \brief mark the end of an old scope.
+   * \param scope_id The scope id to be ended.
+   */
+  void EndScope(int scope_id) {
+    scope_valid_[scope_id] = false;
+    indent_ -= 2;
+  }
+  /*!
+   * \brief Print the indent to the stream.
+   * \param more_indent More indentation besides the current one.
+   */
+  void PrintIndent(int64_t more_indent = 0) {
+    for (int i = 0; i < indent_ + more_indent; ++i) {
+      stream_ << ' ';
+    }
+  }
+  /*!
+   * \brief print end of the line.
+   */
+  void PrintEndInst(const char* suffix) {
+    stream_ << suffix;
+  }
+  /*!
+   * \brief Allocate temporary value
+   * \return A new text value.
+   */
+  TextValue AllocTempVar() {
+    std::ostringstream os;
+    os << '%' << temp_var_counter_++;
+    return TextValue(os.str());
+  }
+  /*!
+   * \brief Allocate name to a variable.
+   * \param var The input variable.
+   * \return The corresponding name.
+   */
+  TextValue AllocVarName(const Var& var) {
+    std::string name = GetUniqueName('%' + var->name_hint);
+    TextValue val(name);
+    CHECK(!memo_.count(var));
+    memo_[var] = val;
+    return val;
+  }
+
+ private:
+  class AttrPrinter;
+  friend class AttrPrinter;
+  /*! \brief meta data context */
+  TextMetaDataContext meta_;
+  /*! \brief Check whether scope is still valid */
+  std::vector<bool> scope_valid_;
+  /*! \brief The current indentation value */
+  int indent_{0};
+  /*! \brief name allocation map */
+  std::unordered_map<std::string, int> name_alloc_map_;
+  /*! \brief Map from expression to its text value */
+  std::unordered_map<Expr, TextValue, NodeHash, NodeEqual> memo_;
+  /*! \brief counter of temporary variable */
+  int64_t temp_var_counter_{0};
+  /*! \brief Output stream */
+  std::ostringstream stream_;
+};
+
+/*!
+ * \brief Attribute printer which prints the attributes in the call.
+ */
+class TextPrinter::AttrPrinter: public AttrVisitor {
+ public:
+  AttrPrinter(std::ostream& stream, TextPrinter* parent)  // NOLINT(*)
+      : stream_(stream), parent_(parent) {}
+
+  void Visit(const char* key, double* value) final {
+    PrintSep();
+    stream_ << key << "=" << value[0];
+  }
+  void Visit(const char* key, int64_t* value) final {
+    PrintSep();
+    stream_ << key << "=" << value[0];
+  }
+  void Visit(const char* key, uint64_t* value) final {
+    PrintSep();
+    stream_ << key << "=" << value[0];
+  }
+  void Visit(const char* key, int* value) final {
+    PrintSep();
+    stream_ << key << "=" << value[0];
+  }
+  void Visit(const char* key, bool* value) final {
+    PrintSep();
+    stream_ << key << "=";
+    parent_->PrintBool(value[0], stream_);
+  }
+  void Visit(const char* key, std::string* value) final {
+    PrintSep();
+    stream_ << key << "=";
+    parent_->PrintString(value[0], stream_);
+  }
+  void Visit(const char* key, void** value) final {
+    LOG(FATAL) << "do not allow void as argument";
+  }
+  void Visit(const char* key, DataType* value) final {
+    PrintSep();
+    stream_ << key << "=";
+    parent_->PrintString(runtime::TVMType2String(Type2TVMType(value[0])), stream_);
+  }
+  void Visit(const char* key, NodeRef* value) final {
+    PrintSep();
+    stream_ << key << "=";
+    parent_->PrintAttr(value[0], stream_);
+  }
+  void Visit(const char* key, runtime::NDArray* value) final {
+    LOG(FATAL) << "do not allow NDarray as argument";
+  }
+
+ private:
+  void PrintSep() {
+    stream_ << ", ";
+  }
+  std::ostream& stream_;  // NOLINT(*)
+  TextPrinter* parent_;
+};
+
+void TextPrinter::PrintCallAttrs(const Expr& op,
+                                 const Attrs& attrs,
+                                 std::ostream& os) {  // NOLINT(*)
+  if (!attrs.defined()) return;
+  if (const auto* op_node = op.as<OpNode>()) {
+    if (attrs->type_index() == op_node->attrs_type_index) {
+      AttrPrinter printer(os, this);
+      const_cast<BaseAttrsNode*>(attrs.operator->())
+          ->VisitNonDefaultAttrs(&printer);
+      return;
+    }
+  }
+  os << ", " << meta_.GetMetaNode(attrs);
+}
+
+std::string RelayPrint(const NodeRef& node) {
+  return TextPrinter().Print(node);
+}
+
+TVM_REGISTER_API("relay._expr._text_print")
+.set_body_typed<std::string(const NodeRef&)>(RelayPrint);
+
+}  // namespace relay
+}  // namespace tvm

src/relay/ir/type.cc

@@ -22,6 +22,8 @@ TensorType TensorTypeNode::Scalar(DataType dtype) {
   return TensorTypeNode::make({}, dtype);
 }
 
+TVM_REGISTER_NODE_TYPE(TensorTypeNode);
+
 TVM_REGISTER_API("relay._make.TensorType")
 .set_body([](TVMArgs args, TVMRetValue *ret) {
   Array<IndexExpr> shape = args[0];
@@ -30,8 +32,8 @@ TVM_REGISTER_API("relay._make.TensorType")
 
 TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable)
 .set_dispatch<TensorTypeNode>([](const TensorTypeNode *node,
-                                     tvm::IRPrinter *p) {
-  p->stream << "TensorTypeNode(" << node->dtype << ", " << node->shape << ")";
+                                 tvm::IRPrinter *p) {
+  p->stream << "TensorType(" << node->shape << ", " << node->dtype << ")";
 });
 
 TypeParam TypeParamNode::make(std::string name, TypeParamNode::Kind kind) {
@@ -41,6 +43,8 @@ TypeParam TypeParamNode::make(std::string name, TypeParamNode::Kind kind) {
   return TypeParam(n);
 }
 
+TVM_REGISTER_NODE_TYPE(TypeParamNode);
+
 TVM_REGISTER_API("relay._make.TypeParam")
 .set_body([](TVMArgs args, TVMRetValue *ret) {
   int kind = args[1];
@@ -61,6 +65,8 @@ IncompleteType IncompleteTypeNode::make(TypeParamNode::Kind kind) {
   return IncompleteType(n);
 }
 
+TVM_REGISTER_NODE_TYPE(IncompleteTypeNode);
+
 TVM_REGISTER_API("relay._make.IncompleteType")
 .set_body([](TVMArgs args, TVMRetValue* ret) {
     int kind = args[0];
@@ -86,6 +92,8 @@ FuncType FuncTypeNode::make(tvm::Array<Type> arg_types,
   return FuncType(n);
 }
 
+TVM_REGISTER_NODE_TYPE(FuncTypeNode);
+
 TVM_REGISTER_API("relay._make.FuncType")
 .set_body([](TVMArgs args, TVMRetValue *ret) {
   *ret = FuncTypeNode::make(args[0], args[1], args[2], args[3]);
@@ -111,6 +119,8 @@ TypeRelation TypeRelationNode::make(TypeRelationFn func,
   return TypeRelation(n);
 }
 
+TVM_REGISTER_NODE_TYPE(TypeRelationNode);
+
 TVM_REGISTER_API("relay._make.TypeRelation")
 .set_body([](TVMArgs args, TVMRetValue *ret) {
     *ret = TypeRelationNode::make(args[0], args[1], args[2], args[3]);
@@ -129,6 +139,8 @@ TupleType TupleTypeNode::make(Array<Type> fields) {
   return TupleType(n);
 }
 
+TVM_REGISTER_NODE_TYPE(TupleTypeNode);
+
 TVM_REGISTER_API("relay._make.TupleType")
 .set_body([](TVMArgs args, TVMRetValue *ret) {
     *ret = TupleTypeNode::make(args[0]);

src/relay/op/nn/convolution.cc

@@ -11,7 +11,7 @@
 namespace tvm {
 namespace relay {
 
-TVM_REGISTER_NODE_TYPE(ConvAttrs);
+TVM_REGISTER_NODE_TYPE(Conv2DAttrs);
 
 bool Conv2DRel(const Array<Type>& types,
                int num_inputs,
@@ -25,7 +25,7 @@ bool Conv2DRel(const Array<Type>& types,
   static const Layout kNCHW("NCHW");
   static const Layout kOIHW("OIHW");
 
-  const ConvAttrs* param = attrs.as<ConvAttrs>();
+  const Conv2DAttrs* param = attrs.as<Conv2DAttrs>();
   CHECK(param != nullptr);
   const Layout in_layout(param->data_layout);
   const Layout kernel_layout(param->weight_layout);
@@ -113,7 +113,7 @@ Expr MakeConv2D(Expr data,
                 std::string weight_layout,
                 std::string out_layout,
                 DataType out_dtype) {
-  auto attrs = make_node<ConvAttrs>();
+  auto attrs = make_node<Conv2DAttrs>();
   attrs->strides = std::move(strides);
   attrs->padding = std::move(padding);
   attrs->dilation = std::move(dilation);
@@ -148,6 +148,7 @@ with the layer input to produce a tensor of outputs.
             (batch_size, channels, out_height, out_width) if `layout` is `NCHW`.
 
 )code" TVM_ADD_FILELINE)
+.set_attrs_type_key("relay.attrs.Conv2DAttrs")
 .set_num_inputs(2)
 .add_argument("data", "Tensor", "The input tensor.")
 .add_argument("weight", "Tensor", "The weight tensor.")
@@ -296,6 +297,7 @@ v            (batch_size, channels, out_height, out_width) if `layout` is `NCHW`
                 out_width = (width-1)*strides[1]-2*padding[1]+kernel_size[1]+output_padding[1]
 
 )code" TVM_ADD_FILELINE)
+.set_attrs_type_key("relay.attrs.Conv2DTransposeAttrs")
 .set_num_inputs(2)
 .add_argument("data", "Tensor", "The input tensor.")
 .add_argument("weight", "Tensor", "The weight tensor.")

src/relay/op/nn/nn.cc

@@ -78,6 +78,7 @@ RELAY_REGISTER_OP("nn.dense")
 - **out**: `(x1, x2, ..., xn, units)`.
 
 )code" TVM_ADD_FILELINE)
+.set_attrs_type_key("relay.attrs.DenseAttrs")
 .set_num_inputs(2)
 .add_argument("data", "nD Tensor", "Input data.")
 .add_argument("weight", "2D Tensor", "Weight matrix.")
@@ -107,6 +108,7 @@ RELAY_REGISTER_OP("nn.leaky_relu")
 `y = x > 0 ? x : alpha * x`
 
 )code" TVM_ADD_FILELINE)
+.set_attrs_type_key("relay.attrs.LeakyReluAttrs")
 .set_num_inputs(1)
 .add_argument("data", "Tensor", "Input data.")
 .set_support_level(3)
@@ -135,6 +137,7 @@ RELAY_REGISTER_OP("nn.softmax")
 
 - **data**: The input data
 )code" TVM_ADD_FILELINE)
+.set_attrs_type_key("relay.attrs.SoftmaxAttrs")
 .set_num_inputs(1)
 .add_argument("data", "Tensor", "The input tensor.")
 .set_support_level(1)
@@ -163,6 +166,7 @@ RELAY_REGISTER_OP("nn.log_softmax")
 
 - **data**: The input data
 )code" TVM_ADD_FILELINE)
+.set_attrs_type_key("relay.attrs.SoftmaxAttrs")
 .set_num_inputs(1)
 .add_argument("data", "Tensor", "The input tensor.")
 .set_support_level(1)
@@ -171,9 +175,9 @@ RELAY_REGISTER_OP("nn.log_softmax")
 
 // BatchFlatten
 bool BatchFlattenRel(const Array<Type>& types,
-               int num_inputs,
-               const Attrs& attrs,
-               const TypeReporter& reporter) {
+                     int num_inputs,
+                     const Attrs& attrs,
+                     const TypeReporter& reporter) {
   CHECK_EQ(types.size(), 2);
   const auto* data = types[0].as<TensorTypeNode>();
   if (data == nullptr) return false;
@@ -278,6 +282,7 @@ centered at that value (zero padding is added where necessary).
 
 - **data**: The input tensor.
 )code" TVM_ADD_FILELINE)
+.set_attrs_type_key("relay.attrs.LRNAttrs")
 .set_num_inputs(1)
 .add_argument("data", "Tensor", "The input tensor.")
 .set_support_level(2)
@@ -296,12 +301,12 @@ Expr MakeL2Normalize(Expr data,
 }
 
 TVM_REGISTER_API("relay.op.nn._make.l2_normalize")
-  .set_body([](const TVMArgs& args, TVMRetValue* rv) {
-      runtime::detail::unpack_call<Expr, 3>(MakeL2Normalize, args, rv);
+.set_body([](const TVMArgs& args, TVMRetValue* rv) {
+    runtime::detail::unpack_call<Expr, 3>(MakeL2Normalize, args, rv);
   });
 
 RELAY_REGISTER_OP("nn.l2_normalize")
-    .describe(R"code(L2 Normalization layer.
+.describe(R"code(L2 Normalization layer.
 
 Normalizes along dimension axis using an L2 norm
 
@@ -352,6 +357,7 @@ During training, each element of the input is set to zero with probability ``p``
 The whole array is rescaled by ``1/(1-p)`` to keep the expected sum of the input unchanged.
 
 )code" TVM_ADD_FILELINE)
+.set_attrs_type_key("relay.attrs.DropoutAttrs")
 .set_num_inputs(1)
 .add_argument("data", "Tensor", "Input to which dropout will be applied.")
 .set_support_level(1)
@@ -478,6 +484,7 @@ axis to be the last item in the input shape.
 .. note::
     This operator can be optimized away for inference.
 )code" TVM_ADD_FILELINE)
+.set_attrs_type_key("relay.attrs.BatchNormAttrs")
 .set_num_inputs(5)
 .add_argument("data", "Tensor", "Input to which batch_norm will be applied.")
 .add_argument("gamma", "Tensor", "The gamma scale factor.")

src/relay/op/nn/pad.cc

@@ -60,7 +60,7 @@ bool PadRel(const Array<Type>& types,
 }
 
 // Handler to create a call to the padding op used by front-end FFI
-  Expr MakePad(Expr data, Array<Array<IndexExpr> > pad_width, double pad_value) {
+Expr MakePad(Expr data, Array<Array<IndexExpr> > pad_width, double pad_value) {
   auto attrs = make_node<PadAttrs>();
   attrs->pad_value = pad_value;
   attrs->pad_width = std::move(pad_width);

src/relay/op/tensor/transform.cc

@@ -76,6 +76,7 @@ RELAY_REGISTER_OP("expand_dims")
 
 )code" TVM_ADD_FILELINE)
 .set_num_inputs(1)
+.set_attrs_type_key("relay.attrs.ExpandDimsAttrs")
 .add_argument("data", "Tensor", "The input tensor.")
 .set_support_level(1)
 .add_type_rel("ExpandDims", ExpandDimsRel);
@@ -481,6 +482,7 @@ RELAY_REGISTER_OP("zeros")
 .describe(R"code(Fill array with zeros.
 
 )code" TVM_ADD_FILELINE)
+.set_attrs_type_key("relay.attrs.InitOpAttrs")
 .set_num_inputs(0)
 .set_support_level(3)
 .add_type_rel("InitOp", InitOpRel);
@@ -503,6 +505,7 @@ RELAY_REGISTER_OP("ones")
 .describe(R"code(Fill array with ones.
 
 )code" TVM_ADD_FILELINE)
+.set_attrs_type_key("relay.attrs.InitOpAttrs")
 .set_num_inputs(0)
 .set_support_level(3)
 .add_type_rel("InitOp", InitOpRel);
@@ -697,6 +700,7 @@ RELAY_REGISTER_OP("squeeze")
 
 )code" TVM_ADD_FILELINE)
 .set_num_inputs(1)
+.set_attrs_type_key("relay.attrs.SqueezeAttrs")
 .add_argument("data", "Tensor", "The input tensor.")
 .set_support_level(3)
 .add_type_rel("Squeeze", SqueezeRel);

src/relay/pass/dead_code.cc

@@ -74,7 +74,10 @@ class CalcDep : private ExprMutator {
   }
 
   Expr VisitExpr_(const FunctionNode* f) final {
-    return FunctionNode::make(f->params, f->ret_type, Eliminate(f->body), f->type_params);
+    return FunctionNode::make(f->params,
+                              Eliminate(f->body),
+                              f->ret_type,
+                              f->type_params);
   }
 
   // generate the let list from dependency graph

src/relay/pass/type_functor.h

@@ -20,6 +20,7 @@ class TypeFunctor;
 #define TYPE_FUNCTOR_DEFAULT \
   { return VisitTypeDefault_(op, std::forward<Args>(args)...); }
 
+
 #define RELAY_TYPE_FUNCTOR_DISPATCH(OP)                                   \
   vtable.template set_dispatch<OP>(                                       \
       [](const NodeRef& n, TSelf* self, Args... args) {                   \

tests/python/relay/test_ir_debug_printer.py

-import tvm
-from tvm import relay
-from tvm.relay.expr import debug_print
-from tvm.relay.ir_builder import IRBuilder
-
-ib = IRBuilder()
-
-def show(e):
-    r = debug_print(ib.env, e)
-    assert r is not None
-
-
-def test_constant():
-    arr = tvm.nd.array(10)
-    const = relay.Constant(arr)
-    show(const)
-    # should print the array inside?
-
-
-def test_tuple():
-    fields = tvm.convert([])
-    tup = relay.Tuple(fields)
-    show(tup)
-
-
-def test_local_var():
-    name_hint = 's'
-    lv = relay.Var(name_hint)
-    show(lv)
-
-
-def test_dup_var():
-    lv = relay.Var('s')
-    rv = relay.Var('s')
-    show(relay.Tuple([lv, rv]))
-
-
-def test_large_dup_var():
-    av = relay.Var('s')
-    bv = relay.Var('s')
-    cv = relay.Var('s')
-    show(relay.Tuple([av, bv, cv]))
-
-
-def test_global_var():
-    name_hint = 'g'
-    gv = relay.GlobalVar(name_hint)
-    gv.name_hint == name_hint
-    show(gv)
-
-
-def test_function():
-    param_names = ['a', 'b', 'c', 'd']
-    params = tvm.convert([relay.Var(n) for n in param_names])
-    ret_type = None
-    body = params[0]
-    type_params = tvm.convert([])
-    fn = relay.Function(params, ret_type, body, type_params)
-    show(fn)
-
-
-
-def test_call():
-    op = relay.Var('f')
-    arg_names = ['a', 'b', 'c', 'd']
-    args = tvm.convert([relay.Var(n) for n in arg_names])
-    call = relay.Call(op, args, None, None)
-    show(call)
-
-
-def test_let():
-    ty = relay.ty.TensorType((10, 20), 'float32')
-    lv = relay.Var('x', ty)
-    arr = tvm.nd.array(10)
-    value = relay.Constant(arr)
-    let = relay.Let(lv, value, lv)
-    show(let)
-
-
-def test_if():
-    cond = relay.Var('cond')
-    left = relay.Var('left')
-    right = relay.Var('right')
-    ife = relay.If(cond, left, right)
-    show(ife)
-
-def test_tuple_get_item():
-    t = relay.Var('t')
-    g = relay.TupleGetItem(t, 0)
-    show(g)

tests/python/relay/test_ir_text_printer.py

+import tvm
+import numpy as np
+from tvm import relay
+
+do_print = [False]
+
+def show(text):
+    if do_print[0]:
+        print("---------------------------")
+        print(text)
+
+def test_func():
+    x = relay.var("x", shape=(3, 2))
+    y = relay.var("y")
+    one = relay.const(10e10, dtype="float32")
+    z = relay.add(x, one)
+    z = relay.add(z, z)
+    f = relay.Function([x, y], z)
+    show(z.astext())
+    show(f.astext())
+
+
+def test_env():
+    x = relay.var("x", "float32")
+    y = relay.var("y", "float32")
+    z = relay.add(x, y)
+    z = relay.add(z, z)
+    f = relay.Function([x, y], z)
+    env = relay.Environment()
+    env.add("myf", f)
+    text = env.astext()
+    assert "def @myf" in text
+    assert "%1 = add(%0, %0) # ty=float32" in text
+    show(text)
+
+
+def test_meta_data():
+    n, c, h, w = tvm.var("n"), 10, 224, 224
+    x = relay.var("x", shape=(n, c, h, w))
+    w = relay.var("w")
+    z = relay.nn.conv2d(x, w,
+                        kernel_size=(3, 3),
+                        padding=(1, 1),
+                        channels=2)
+    f = relay.Function([x, w], z)
+    text = f.astext()
+    assert "channels=2" in text
+    assert "meta.Variable(id=0)" in text
+    show(text)
+
+    text = relay.const([1,2,3]).astext()
+    assert "meta.relay.Constant(id=0)" in text
+    show(text)
+
+
+def test_call_attrs():
+    x = relay.var("x")
+    # non default args
+    z = relay.nn.softmax(x, axis=2)
+    assert "axis=2" in z.astext()
+    # default args
+    z = relay.nn.softmax(x)
+    assert "softmax(%x)" in z.astext()
+    # non default args
+    z = relay.expand_dims(x, axis=2, num_newaxis=2)
+    assert "num_newaxis=2" in z.astext()
+
+
+def test_let_if_scope():
+    x = relay.var("x", "float32")
+    y = relay.var("y", "float32")
+    cond = relay.var("cond", "bool")
+    v1 = relay.var("v")
+    v2 = relay.var("v", "float32")
+    then_branch = relay.Let(
+        v1, relay.const(1, "float32"),
+        relay.Let(v2, x, relay.subtract(v1, v2)))
+    v3 = relay.var("v")
+    let2 = relay.Let(v3, y, v3)
+    else_branch = relay.add(let2, let2)
+    result = relay.If(cond, then_branch, else_branch)
+    f = relay.Function([x, y, cond], result)
+    text = f.astext()
+    assert text.count("{") == 4
+    assert "%cond: bool" in text
+    show(f.astext())
+
+
+if __name__ == "__main__":
+    do_print[0] = True
+    test_let_if_scope()
+    test_func()
+    test_env()
+    test_meta_data()
+    test_call_attrs()

tests/python/relay/test_ir_well_formed.py

@@ -10,7 +10,7 @@ def test_well_formed():
     let = relay.Let(x, v, x)
     assert well_formed(let)
     assert not well_formed(relay.Let(x, v, let))
-    f = relay.Function([x], ty, x)
+    f = relay.Function([x], x, ty)
     assert well_formed(f)
     # this test should pass in case of weak uniqueness (only test for shadowing)
     # but we want all binder to be distinct from each other.

tests/python/relay/test_pass_alpha_equal.py

@@ -262,49 +262,49 @@ def test_function_alpha_equal():
     basic_args = [relay.Var("v3", tt1), relay.Var("v4", tt2)]
     basic_tps = [tp1, tp2]
 
-    func = relay.Function([v1, v2],
-                          tt2, v1, basic_tps)
-    mapped = relay.Function(basic_args, tt2, basic_args[0], basic_tps)
+    func = relay.Function([v1, v2], v1,
+                          tt2, basic_tps)
+    mapped = relay.Function(basic_args, basic_args[0], tt2, basic_tps)
     assert alpha_equal(func, mapped)
 
-    fewer_params = relay.Function([relay.Var("v4", tt2)], tt2, v4, basic_tps)
+    fewer_params = relay.Function([relay.Var("v4", tt2)], v4, tt2, basic_tps)
     assert not alpha_equal(func, fewer_params)
 
     more_params = relay.Function([relay.Var("v3", tt1),
                                   relay.Var("v4", tt2),
-                                  relay.Var("v2", tt2)], tt2, v4, basic_tps)
+                                  relay.Var("v2", tt2)], v4, tt2, basic_tps)
     assert not alpha_equal(func, more_params)
 
-    params_unordered = relay.Function([v2, v1],
-                                      tt2, v1, basic_tps)
+    params_unordered = relay.Function([v2, v1], v1,
+                                      tt2, basic_tps)
     assert not alpha_equal(func, params_unordered)
 
-    params_mismatch = relay.Function([v1, v3],
-                                     tt2, v1, basic_tps)
+    params_mismatch = relay.Function([v1, v3], v1,
+                                     tt2, basic_tps)
     assert not alpha_equal(func, params_mismatch)
 
     # also would not typecheck
-    ret_type_mismatch = relay.Function(basic_args, tt1, v4, basic_tps)
+    ret_type_mismatch = relay.Function(basic_args, v4, tt1, basic_tps)
     assert not alpha_equal(func, ret_type_mismatch)
 
     # also mis-typed
-    different_body = relay.Function(basic_args, tt2, v3, basic_tps)
+    different_body = relay.Function(basic_args, v3, tt2, basic_tps)
     assert not alpha_equal(func, different_body)
 
-    fewer_type_params = relay.Function(basic_args, tt2, v4, [tp1])
+    fewer_type_params = relay.Function(basic_args, v4, tt2, [tp1])
     assert not alpha_equal(func, fewer_type_params)
 
-    more_type_params = relay.Function(basic_args, tt2, v4, [tp1, tp2, tp3])
+    more_type_params = relay.Function(basic_args, v4, tt2, [tp1, tp2, tp3])
     assert not alpha_equal(func, more_type_params)
 
-    type_params_unordered = relay.Function(basic_args, tt2, v4, [tp2, tp1])
+    type_params_unordered = relay.Function(basic_args, v4, tt2, [tp2, tp1])
     assert not alpha_equal(func, type_params_unordered)
 
-    different_type_params = relay.Function(basic_args, tt2, v4, [tp3, tp4])
+    different_type_params = relay.Function(basic_args, v4, tt2, [tp3, tp4])
     assert not alpha_equal(func, different_type_params)
 
     # a well-typed example that also differs in body, ret type, and type params
-    tupled_example = relay.Function(basic_args, tt3, relay.Tuple([v3, v4]))
+    tupled_example = relay.Function(basic_args, relay.Tuple([v3, v4]), tt3)
     assert not alpha_equal(func, tupled_example)
 
 

tests/python/relay/test_pass_dead_code_elimination.py

@@ -59,7 +59,7 @@ def test_recursion():
     n = relay.Var("n", e.int32)
     data = relay.Var("data", e.float32)
     funcbody = relay.If(equal(n, convert(0)), data, f(subtract(n, convert(1.0)), log(data)))
-    value = relay.Function([n, data], e.float32, funcbody, [])
+    value = relay.Function([n, data], funcbody, e.float32, [])
     orig = relay.Let(f, funcbody, f(convert(2.0), convert(10000.0)))
     assert alpha_equal(dead_code_elimination(orig), orig)
     assert alpha_equal(dead_code_elimination(relay.Let(f, funcbody, e.three)), e.three)

tests/python/relay/test_pass_free_vars.py

@@ -13,7 +13,7 @@ def test_free_vars():
     let = relay.Let(x, v, x)
     fvx = free_vars(let)
     assert len(free_vars(let)) == 0
-    f = relay.Function([x], ty, x)
+    f = relay.Function([x], x, ty)
     assert len(free_vars(f)) == 0