better type infer ability

apache · Jan 23, 2019 · 0ce1f68 · 0ce1f68
1 parent f5258a8
commit 0ce1f68
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 17 deletions.
diff --git a/src/relay/ir/expr.cc b/src/relay/ir/expr.cc
@@ -315,7 +315,7 @@ TVM_REGISTER_API("relay._make.RefWrite").set_body([](TVMArgs args, TVMRetValue*
 
 TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable)
 .set_dispatch<RefWriteNode>([](const RefWriteNode* node, tvm::IRPrinter* p) {
-  p->stream << "RefWriteNode(" << node->value << ")";
+  p->stream << "RefWriteNode(" << node->ref << ", " << node->value << ")";
 });
 
 TVM_REGISTER_API("relay._expr.TempExprRealize")

diff --git a/src/relay/pass/type_infer.cc b/src/relay/pass/type_infer.cc
@@ -396,25 +396,15 @@ class TypeInferencer : private ExprFunctor<Type(const Expr&)> {
   }
 
   Type VisitExpr_(const RefReadNode* op) final {
-    // TODO(M.K.)
-    // handle case where ref type is not known
-    Type ref_type = GetType(op->ref);
-    auto ref_ty_node = ref_type.as<RefTypeNode>();
-    if (!ref_ty_node) {
-      LOG(FATAL) << "only expressions with ref types is accepted" << GetRef<RefRead>(op);
-    }
-    return ref_ty_node->value;
+    Type it = IncompleteTypeNode::make(TypeVarNode::Kind::kType);
+    this->Unify(GetType(op->ref), RefTypeNode::make(it), op->span);
+    return it;
   }
 
   Type VisitExpr_(const RefWriteNode* op) final {
-    // TODO(M.K.)
-    // handle case where ref type is not known
-    Type ref_type = GetType(op->ref);
-    auto ref_ty_node = ref_type.as<RefTypeNode>();
-    if (!ref_ty_node) {
-      LOG(FATAL) << "only expressions with ref types is accepted" << GetRef<RefWrite>(op);
-    }
-    this->Unify(ref_ty_node->value, GetType(op->value), op->span);
+    Type it = IncompleteTypeNode::make(TypeVarNode::Kind::kType);
+    this->Unify(GetType(op->ref), RefTypeNode::make(it), op->span);
+    this->Unify(GetType(op->value), it, op->span);
     return TupleTypeNode::make({});
   }
 };