Merge pull request #7 from MarisaKirisame/type_of

Type of

Author: jroesch

include/tvm/relay/expr.h

@@ -486,6 +486,37 @@ std::string RelayPrint(
     const NodeRef& node,
     bool show_meta_data = true,
     runtime::TypedPackedFunc<std::string(Expr)> annotate = nullptr);
+
+/*!
+ * \brief User defined type relation, is an input-output relation on types.
+ */
+class TypeOf;
+/*!
+ * \brief TypeRelation container.
+ * \note This node is not directly serializable.
+ * The type function need to be lookedup in the module.
+ */
+class TypeOfNode : public TypeNode {
+ public:
+  /*!
+   * \brief The function on input and output variables which
+   *  this is not directly serializable,
+   *  need to be looked-up in the module.
+   */
+  relay::Expr expr;
+
+  void VisitAttrs(tvm::AttrVisitor* v) final {
+    v->Visit("expr", &expr);
+  }
+
+  TVM_DLL static TypeOf make(relay::Expr expr);
+
+  static constexpr const char* _type_key = "relay.TypeOf";
+  TVM_DECLARE_NODE_TYPE_INFO(TypeOfNode, TypeNode);
+};
+
+RELAY_DEFINE_NODE_REF(TypeOf, TypeOfNode, Type);
+
 }  // namespace relay
 }  // namespace tvm
 #endif  // TVM_RELAY_EXPR_H_

python/tvm/relay/__init__.py

@@ -48,6 +48,7 @@
 scalar_type = ty.scalar_type
 GlobalTypeVar = ty.GlobalTypeVar
 TypeCall = ty.TypeCall
+TypeOf = ty.TypeOf
 
 # Expr
 Expr = expr.Expr

python/tvm/relay/op/_tensor.py

@@ -6,7 +6,7 @@
 from .op import register_gradient
 from .op import schedule_injective, OpPattern
 from .transform import collapse_sum_like
-from .tensor import negative
+from .tensor import negative, zeros_like, ones_like
 
 
 def add_grad(orig, grad):
@@ -29,6 +29,12 @@ def multiply_grad(orig, grad):
 
 register_gradient("multiply", multiply_grad)
 
+def take_grad(orig, grad):
+    return [zeros_like(orig.args[0]), zeros_like(orig.args[1])]
+
+
+register_gradient("take", take_grad)
+
 schedule_broadcast = schedule_injective
 schedule_elemwise = schedule_injective
 

python/tvm/relay/ty.py

@@ -270,6 +270,10 @@ def __init__(self, func, args, num_inputs, attrs):
         self.__init_handle_by_constructor__(_make.TypeRelation,
                                             func, args, num_inputs, attrs)
 
+@register_relay_node
+class TypeOf(Type):
+    def __init__(self, expr):
+        self.__init_handle_by_constructor__(_make.TypeOf, expr)
 
 def scalar_type(dtype):
     """Creates a scalar type.

src/relay/ir/error.cc

@@ -87,9 +87,9 @@ void ErrorReporter::RenderErrors(const Module& module, bool use_color) {
       std::endl <<
       rang::style::reset;
 
-    // for (auto pair : err_map) {
-    //   std::cout << "Key: " << pair.first << " Value: " << pair.second << std::endl;
-    // }
+    for (auto pair : err_map) {
+      std::cout << "Key: " << pair.first << std::endl << " Value: " << pair.second << std::endl;
+    }
 
     // We then call into the Relay printer to generate the program.
     //

src/relay/ir/expr.cc

@@ -283,6 +283,24 @@ TVM_REGISTER_API("relay._expr.TempExprRealize")
     *ret = temp->Realize();
 });
 
+TypeOf TypeOfNode::make(relay::Expr expr) {
+  NodePtr<TypeOfNode> n = make_node<TypeOfNode>();
+  n->expr = std::move(expr);
+  return TypeOf(n);
+}
+
+TVM_REGISTER_NODE_TYPE(TypeOfNode);
+
+TVM_REGISTER_API("relay._make.TypeOf")
+.set_body([](TVMArgs args, TVMRetValue* ret) {
+    *ret = TypeOfNode::make(args[0]);
+});
+
+TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable)
+.set_dispatch<TypeOfNode>([](const TypeOfNode* node,
+                                tvm::IRPrinter* p) {
+  p->stream << "TypeOf(" << node->expr << ")";
+});
 
 }  // namespace relay
 }  // namespace tvm

src/relay/ir/type.cc

@@ -207,5 +207,6 @@ TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable)
   p->stream << "TupleTypeNode(" << node->fields << ")";
 });
 
+
 }  // namespace relay
 }  // namespace tvm

src/relay/ir/type_functor.cc

@@ -68,6 +68,8 @@ void TypeVisitor::VisitType_(const TypeDataNode* op) {
   }
 }
 
+void TypeVisitor::VisitType_(const TypeOfNode* op) {}
+
 // Type Mutator.
 Array<Type> TypeMutator::MutateArray(Array<Type> arr) {
   // The array will do copy on write
@@ -172,6 +174,10 @@ Type TypeMutator::VisitType_(const TypeDataNode* op) {
   return GetRef<Type>(op);
 }
 
+Type TypeMutator::VisitType_(const TypeOfNode* op) {
+  return GetRef<Type>(op);
+}
+
 // Implements bind.
 class TypeBinder : public TypeMutator {
  public:

src/relay/ir/type_functor.h

@@ -72,6 +72,7 @@ class TypeFunctor<R(const Type& n, Args...)> {
   virtual R VisitType_(const GlobalTypeVarNode* op, Args... args) TYPE_FUNCTOR_DEFAULT;
   virtual R VisitType_(const TypeCallNode* op, Args... args) TYPE_FUNCTOR_DEFAULT;
   virtual R VisitType_(const TypeDataNode* op, Args... args) TYPE_FUNCTOR_DEFAULT;
+  virtual R VisitType_(const TypeOfNode* op, Args... args) TYPE_FUNCTOR_DEFAULT;
 
   virtual R VisitTypeDefault_(const Node* op, Args...) {
     LOG(FATAL) << "Do not have a default for " << op->type_key();
@@ -93,6 +94,7 @@ class TypeFunctor<R(const Type& n, Args...)> {
     RELAY_TYPE_FUNCTOR_DISPATCH(GlobalTypeVarNode);
     RELAY_TYPE_FUNCTOR_DISPATCH(TypeCallNode);
     RELAY_TYPE_FUNCTOR_DISPATCH(TypeDataNode);
+    RELAY_TYPE_FUNCTOR_DISPATCH(TypeOfNode);
     return vtable;
   }
 };
@@ -111,6 +113,7 @@ class TypeVisitor : public TypeFunctor<void(const Type& n)> {
   void VisitType_(const GlobalTypeVarNode* op) override;
   void VisitType_(const TypeCallNode* op) override;
   void VisitType_(const TypeDataNode* op) override;
+  void VisitType_(const TypeOfNode* op) override;
 };
 
 // Mutator that transform a type to another one.
@@ -125,6 +128,7 @@ class TypeMutator : public TypeFunctor<Type(const Type& n)> {
   Type VisitType_(const GlobalTypeVarNode* op) override;
   Type VisitType_(const TypeCallNode* op) override;
   Type VisitType_(const TypeDataNode* op) override;
+  Type VisitType_(const TypeOfNode* op) override;
 
  private:
   Array<Type> MutateArray(Array<Type> arr);

src/relay/pass/kind_check.cc

@@ -156,6 +156,10 @@ struct KindChecker : TypeFunctor<Kind(const Type&)> {
     return Kind::kTypeData;
   }
 
+  Kind VisitType_(const TypeOfNode* op) override {
+    return kType;
+  }
+
   Kind Check(const Type& t) {
     return this->VisitType(t);
   }