[Relay] Unifier hotfix (#2437)

include/tvm/relay/pass.h

@@ -108,6 +108,17 @@ bool AlphaEqual(const Type& t1, const Type& t2);
  */
 bool WellFormed(const Expr& expr);
 
+/*! \brief Get all bound variables from expression expr.
+ *
+ * Bound variables are all variables that are declared in the expr.
+ * They only have meaning inside that expr, and can only be used in it.
+ *
+ * \param expr the expression.
+ *
+ * \return List of bound vars, in the PostDFS order in the expression.
+ */
+tvm::Array<Var> BoundVars(const Expr& expr);
+
 /*! \brief Get free type parameters from expression expr.
  *
  * Free variables are variables that are not bound by a
@@ -119,6 +130,14 @@ bool WellFormed(const Expr& expr);
  */
 tvm::Array<Var> FreeVars(const Expr& expr);
 
+/*! \brief Get all variables from expression expr.
+ *
+ * \param expr the expression.
+ *
+ * \return List of all vars, in the PostDFS order in the expression.
+ */
+tvm::Array<Var> AllVars(const Expr& expr);
+
 /*! \brief Get free TypeVars from expression expr.
  *
  * Free type parameters are type parameters that are not bound by a function
@@ -130,6 +149,55 @@ tvm::Array<Var> FreeVars(const Expr& expr);
  */
 tvm::Array<TypeVar> FreeTypeVars(const Expr& expr);
 
+/*! \brief Get free TypeVars from type t.
+ *
+ * Free type parameters are type parameters that are not bound by a function
+ * type in the context.
+ *
+ * \param t the type.
+ *
+ * \return List of free type vars, in the PostDFS order visited by type.
+ */
+tvm::Array<TypeVar> FreeTypeVars(const Type& t);
+
+/*! \brief Get all bound type variables from expression expr.
+ *
+ * Bound variables are all type variables that are declared in the expr.
+ * They only have meaning inside that expr, and can only be used in it.
+ *
+ * \param expr the expression.
+ *
+ * \return List of bound type vars, in the PostDFS order in the expression.
+ */
+tvm::Array<TypeVar> BoundTypeVars(const Expr& expr);
+
+/*! \brief Get all bound type variables from type t.
+ *
+ * Bound variables are all type variables that are declared in the type.
+ * They only have meaning inside that type, and can only be used in it.
+ *
+ * \param t the type
+ *
+ * \return List of bound type vars, in the PostDFS order visited by type.
+ */
+tvm::Array<TypeVar> BoundTypeVars(const Type& t);
+
+/*! \brief Get all type variables in expression expr.
+ *
+ * \param expr the expression.
+ *
+ * \return List of type vars, in the PostDFS order in the expression.
+ */
+tvm::Array<TypeVar> AllTypeVars(const Expr& expr);
+
+/*! \brief Get all type variables in type t.
+ *
+ * \param t the type.
+ *
+ * \return List of type vars, in the PostDFS order visited by type.
+ */
+tvm::Array<TypeVar> AllTypeVars(const Type& t);
+
 /*! \brief Remove expressions which does not effect the program result.
  *
  * It will remove let bindings which are not referenced, and branches that will

python/tvm/relay/ir_pass.py

@@ -158,6 +158,38 @@ def free_vars(expr):
     return _ir_pass.free_vars(expr)
 
 
+def bound_vars(expr):
+    """Get bound vars from expression expr in post-DFS order.
+
+    Parameters
+    ----------
+    expr: tvm.relay.Expr
+        The input expression
+
+    Returns
+    -------
+    free : List[tvm.relay.Var]
+        The list of bound variables in post-DFS order.
+    """
+    return _ir_pass.bound_vars(expr)
+
+
+def all_vars(expr):
+    """Get all vars from expression expr in post-DFS order.
+
+    Parameters
+    ----------
+    expr: tvm.relay.Expr
+        The input expression
+
+    Returns
+    -------
+    free : List[tvm.relay.Var]
+        The list of all variables in post-DFS order.
+    """
+    return _ir_pass.all_vars(expr)
+
+
 def free_type_vars(expr):
     """Get free type variables from expression/type e
 
@@ -168,12 +200,44 @@ def free_type_vars(expr):
 
     Returns
     -------
-    free : List[tvm.relay.TypeParam]
-        The list of free type variables
+    free : List[tvm.relay.TypeVar]
+        The list of free type variables in post-DFS order
     """
     return _ir_pass.free_type_vars(expr)
 
 
+def bound_type_vars(expr):
+    """Get bound type variables from expression/type e
+
+    Parameters
+    ----------
+    expr: Union[tvm.relay.Expr,tvm.relay.Type]
+        The input expression/type
+
+    Returns
+    -------
+    free : List[tvm.relay.TypeVar]
+        The list of bound type variables in post-DFS order
+    """
+    return _ir_pass.bound_type_vars(expr)
+
+
+def all_type_vars(expr):
+    """Get all type variables from expression/type e
+
+    Parameters
+    ----------
+    expr: Union[tvm.relay.Expr,tvm.relay.Type]
+        The input expression/type
+
+    Returns
+    -------
+    free : List[tvm.relay.TypeVar]
+        The list of all type variables in post-DFS order
+    """
+    return _ir_pass.all_type_vars(expr)
+
+
 def simplify_inference(expr):
     """ Simplify the data-flow graph for inference phase.
 

src/relay/pass/gradient.cc

@@ -205,14 +205,25 @@ Expr FirstOrderGradient(const Expr& re, const Module& mod) {
       });
     return Pair(res.foward, grad);
   });
+
+  // if type annotations are provided, we will construct a ret type;
+  // otherwise, leave it to be inferred
+  Type ret_type = Type();
   std::vector<Type> vt;
+  bool missing = !f->ret_type.defined();
   for (const auto& p : f->params) {
+    if (missing || !p->type_annotation.defined()) {
+      missing = true;
+      break;
+    }
     vt.push_back(p->type_annotation);
   }
-  return FunctionNode::make(f->params,
-                            body,
-                            TupleTypeNode::make({f->ret_type, TupleTypeNode::make({})}),
-                            {});
+
+  if (!missing) {
+    ret_type = TupleTypeNode::make({f->ret_type, TupleTypeNode::make(vt)});
+  }
+
+  return FunctionNode::make(f->params, body, ret_type, {});
 }
 
 TVM_REGISTER_API("relay._ir_pass.first_order_gradient")

src/relay/pass/type_infer.cc

@@ -56,31 +56,11 @@ bool TupleGetItemRel(const Array<Type>& types,
   return true;
 }
 
-bool MakeTupleRel(const Array<Type>& types,
-                  int num_inputs,
-                  const Attrs& attrs,
-                  const TypeReporter& reporter) {
-  CHECK_EQ(static_cast<size_t>(num_inputs + 1), types.size());
-  for (int i = 0; i < num_inputs; ++i) {
-    if (types[i].as<IncompleteTypeNode>()) return false;
-  }
-  Array<Type> fields;
-  for (int i = 0; i < num_inputs; ++i) {
-    fields.push_back(types[i]);
-  }
-  reporter->Assign(types[num_inputs], TupleTypeNode::make(fields));
-  return true;
-}
-
 TVM_REGISTER_NODE_TYPE(TupleGetItemAttrs);
 TVM_REGISTER_API("tvm.relay.type_relation.TupleGetItem")
 .set_body_typed<bool(const Array<Type>&, int, const Attrs&, const TypeReporter&)>(
     TupleGetItemRel);
 
-TVM_REGISTER_API("tvm.relay.type_relation.MakeTuple")
-.set_body_typed<bool(const Array<Type>&, int, const Attrs&, const TypeReporter&)>(
-    MakeTupleRel);
-
 struct ResolvedTypeInfo {
   explicit ResolvedTypeInfo(Type checked_type, Array<Type> type_args)
       : checked_type(checked_type), type_args(type_args) {}
@@ -120,6 +100,10 @@ class TypeInferencer : private ExprFunctor<Type(const Expr&)> {
   // type inferencer will populate it up
   std::unordered_map<Expr, ResolvedTypeInfo, NodeHash, NodeEqual> type_map_;
 
+  // used to ensure we don't have free type vars hanging around
+  // (a temporary measure until we have proper generalization implemented)
+  Map<TypeVar, Type> instantiation_map_;
+
   // The solver used by the inferencer.
   TypeSolver solver_;
   // relation function
@@ -140,6 +124,32 @@ class TypeInferencer : private ExprFunctor<Type(const Expr&)> {
       return Type();
     }
   }
+
+  // Substitutes every type var in t with a corresponding incomplete type.
+  // This is a temporary measure to ensure type vars behave until
+  // generalization is properly implemented.
+  Type Instantiate(const Type &t) {
+    if (!t.defined()) {
+      return t;
+    }
+    auto* ft = t.as<FuncTypeNode>();
+    if (ft == nullptr) {
+      return Bind(t, instantiation_map_);
+    }
+
+    for (auto type_param : ft->type_params) {
+      instantiation_map_.Set(type_param, IncompleteTypeNode::make(TypeVarNode::Kind::kType));
+    }
+
+    Type ret_type = ft->ret_type;
+    if (!ret_type.defined()) {
+      ret_type = IncompleteTypeNode::make(TypeVarNode::Kind::kType);
+    }
+
+    auto strip_tvs = FuncTypeNode::make(ft->arg_types, ret_type, {}, ft->type_constraints);
+    return Bind(strip_tvs, instantiation_map_);
+  }
+
   // Lazily get type for expr
   // will call visit to deduce it if it is not in the type_map_
   Type GetType(const Expr &expr) {
@@ -147,7 +157,7 @@ class TypeInferencer : private ExprFunctor<Type(const Expr&)> {
     if (it != type_map_.end() && it->second.checked_type.defined()) {
       return it->second.checked_type;
     }
-    Type ret = this->VisitExpr(expr);
+    Type ret = Instantiate(this->VisitExpr(expr));
     ResolvedTypeInfo& rti = type_map_[expr];
     rti.checked_type = ret;
     return ret;
@@ -175,19 +185,11 @@ class TypeInferencer : private ExprFunctor<Type(const Expr&)> {
   }
 
   Type VisitExpr_(const TupleNode* op) final {
-    if (!make_tuple_rel_.defined())  {
-      make_tuple_rel_ = TypeRelationFn(
-          EnvFunc::Get("tvm.relay.type_relation.MakeTuple").node_);
-    }
     Array<Type> types;
     for (Expr field : op->fields) {
       types.push_back(GetType(field));
     }
-    Type rtype = IncompleteTypeNode::make(TypeVarNode::Kind::kType);
-    types.push_back(rtype);
-    solver_.AddConstraint(TypeRelationNode::make(
-        make_tuple_rel_, types, op->fields.size(), Attrs()));
-    return rtype;
+    return TupleTypeNode::make(types);
   }
 
   Type VisitExpr_(const TupleGetItemNode* op) final {
@@ -209,11 +211,17 @@ class TypeInferencer : private ExprFunctor<Type(const Expr&)> {
   }
 
   Type VisitExpr_(const LetNode* op) final {
+    // if the definition is a function literal, permit recursion
+    bool is_functional_literal = op->value.as<FunctionNode>() != nullptr;
+    if (is_functional_literal) {
+      type_map_[op->var].checked_type = IncompleteTypeNode::make(TypeVarNode::Kind::kType);
+    }
+
     Type vtype = GetType(op->value);
     if (op->var->type_annotation.defined()) {
       vtype = Unify(vtype, op->var->type_annotation, op->span);
     }
-    CHECK(!type_map_.count(op->var));
+    CHECK(is_functional_literal || !type_map_.count(op->var));
     // NOTE: no scoping is necessary because var are unique in program
     type_map_[op->var].checked_type = vtype;
     return GetType(op->body);
@@ -252,16 +260,14 @@ class TypeInferencer : private ExprFunctor<Type(const Expr&)> {
     return rtype;
   }
 
-  // instantiate the function type with fresh
-  FuncType Instantiate(const FuncTypeNode* fn_ty, Array<Type>* ty_args) {
+  // substitute the type args in the function type
+  FuncType InstantiateFuncType(const FuncTypeNode* fn_ty, const Array<Type>& ty_args) {
     tvm::Map<TypeVar, Type> subst_map;
 
     // Build a subsitituion map up from the function type and type arguments.
     // Eventually allow the type vars to be passed in.
-    for (auto ty_param : fn_ty->type_params) {
-      IncompleteType fresh = IncompleteTypeNode::make(ty_param->kind);
-      subst_map.Set(ty_param, fresh);
-      ty_args->push_back(fresh);
+    for (size_t i = 0; i < fn_ty->type_params.size(); i++) {
+      subst_map.Set(fn_ty->type_params[i], ty_args[i]);
     }
 
     Type ret_type = fn_ty->ret_type;
@@ -296,13 +302,32 @@ class TypeInferencer : private ExprFunctor<Type(const Expr&)> {
   Type GeneralCall(const CallNode* call, Array<Type> arg_types) {
     Type ftype = GetType(call->op);
     auto* fn_ty_node = ftype.as<FuncTypeNode>();
+    auto* inc_ty_node = ftype.as<IncompleteTypeNode>();
+
+    CHECK(fn_ty_node != nullptr || inc_ty_node != nullptr)
+      << "only expressions with function types can be called, found "
+      << ftype << " at " << call->span;
+
+    // incomplete type => it must be a function taking the arg types
+    // with an unknown return type
+    if (inc_ty_node != nullptr) {
+      Type ret_type = IncompleteTypeNode::make(TypeVarNode::Kind::kType);
+      Type func_type = FuncTypeNode::make(arg_types, ret_type, {}, {});
+      Type unified = this->Unify(ftype, func_type, call->span);
+      fn_ty_node = unified.as<FuncTypeNode>();
+    }
 
-    CHECK(fn_ty_node != nullptr)
-        << "only expressions with function types can be called, found "
-        << ftype << " at " << call->span;
-
-    Array<Type> type_args;
-    FuncType fn_ty = Instantiate(fn_ty_node, &type_args);
+    Array<Type> type_args = call->type_args;
+    if (type_args.size() == 0) {
+      for (size_t i = 0; i < fn_ty_node->type_params.size(); i++) {
+        type_args.push_back(IncompleteTypeNode::make(TypeVarNode::Kind::kType));
+      }
+    }
+    CHECK(type_args.size() == fn_ty_node->type_params.size())
+      << "Incorrect number of type args in " << call->span << ": "
+      << "Expected " << fn_ty_node->type_params.size()
+      << "but got " << type_args.size();
+    FuncType fn_ty = InstantiateFuncType(fn_ty_node, type_args);
 
     AddTypeArgs(GetRef<Call>(call), type_args);
 
@@ -353,26 +378,17 @@ class TypeInferencer : private ExprFunctor<Type(const Expr&)> {
   }
 
   Type VisitExpr_(const FunctionNode* f) final {
+    solver_.Solve();
+    Array<Type> arg_types;
     for (auto param : f->params) {
-      GetType(param);
+      arg_types.push_back(GetType(param));
     }
     Type rtype = GetType(f->body);
-    // Run solver using the currently known information
-    solver_.Solve();
-    // Trying to resolve
-    Array<Type> arg_types;
-    for (size_t i = 0; i < f->params.size(); ++i) {
-      Type atype = solver_.Resolve(GetType(f->params[i]));
-      CHECK(atype.as<IncompleteTypeNode>() == nullptr)
-          << "Cannot resolve type of " << i
-          << "-th parameter of function at" << f->span;
-      arg_types.push_back(atype);
+    if (f->ret_type.defined()) {
+      rtype = this->Unify(f->ret_type, rtype, f->span);
     }
-    rtype = solver_.Resolve(rtype);
-    CHECK(rtype.as<IncompleteTypeNode>() == nullptr)
-        << "Cannot resolve return type of function at" << f->span;
-    // do not support constraint lifting for now.
-    return FuncTypeNode::make(arg_types, rtype, f->type_params, {});
+    auto ret = FuncTypeNode::make(arg_types, rtype, f->type_params, {});
+    return solver_.Resolve(ret);
   }
 };
 
@@ -380,7 +396,7 @@ class TypeInferencer::Resolver : public ExprMutator {
  public:
   Resolver(const std::unordered_map<Expr, ResolvedTypeInfo, NodeHash, NodeEqual>& tmap,
            TypeSolver* solver)
-      : tmap_(tmap), solver_(solver) {
+    : tmap_(tmap), solver_(solver) {
   }
 
   Expr VisitExpr_(const VarNode* op) final {
@@ -525,6 +541,7 @@ Expr TypeInferencer::Infer(Expr expr) {
   GetType(expr);
   // Step 1: Solve the constraints.
   solver_.Solve();
+
   // Step 2: Attach resolved types to checked_type field.
   auto resolved_expr = Resolver(type_map_, &solver_).VisitExpr(expr);
   CHECK(WellFormed(resolved_expr));

src/relay/pass/type_solver.cc

@@ -5,6 +5,7 @@
  */
 #include <string>
 #include "type_solver.h"
+#include "../ir/type_functor.h"
 
 namespace tvm {
 namespace relay {
@@ -38,9 +39,298 @@ class TypeSolver::Reporter : public TypeReporterNode {
   TypeSolver* solver_;
 };
 
+class TypeSolver::OccursChecker : public TypeVisitor {
+ public:
+  explicit OccursChecker(TypeSolver* solver, TypeNode* var)
+    : solver_(solver), var_(var), found_(false) {}
+
+  bool Check(const Type& t) {
+    VisitType(t);
+    return found_;
+  }
+
+  void VisitType_(const IncompleteTypeNode* op) override {
+    IncompleteType t = GetRef<IncompleteType>(op);
+    TypeNode* node = solver_->GetTypeNode(t);
+    found_ = found_ || (var_->FindRoot() == node->FindRoot());
+  }
+
+ private:
+  TypeSolver* solver_;
+  TypeNode* var_;
+  bool found_;
+};
+
+class TypeSolver::Unifier : public TypeFunctor<Type(const Type&, const Type&)> {
+ public:
+  explicit Unifier(TypeSolver* solver) : solver_(solver) {}
+
+  Type Unify(const Type& src, const Type& dst) {
+    // Known limitation
+    // - handle shape pattern matching
+    TypeNode* lhs = solver_->GetTypeNode(dst);
+    TypeNode* rhs = solver_->GetTypeNode(src);
+
+    // do occur check so we don't create self-referencing structure
+    if (lhs->FindRoot() == rhs->FindRoot()) {
+      return lhs->resolved_type;
+    }
+    if (lhs->resolved_type.as<IncompleteTypeNode>()) {
+      CHECK(!CheckOccurs(lhs, rhs->resolved_type))
+        << "Incomplete type " << lhs->resolved_type << " occurs in "
+        << rhs->resolved_type << ", cannot unify";
+      solver_->MergeFromTo(lhs, rhs);
+      return rhs->resolved_type;
+    } else if (rhs->resolved_type.as<IncompleteTypeNode>()) {
+      CHECK(!CheckOccurs(rhs, lhs->resolved_type))
+        << "Incomplete type " << rhs->resolved_type << " occurs in "
+        << lhs->resolved_type << ", cannot unify";
+      solver_->MergeFromTo(rhs, lhs);
+      return lhs->resolved_type;
+    } else {
+      Type resolved = this->VisitType(lhs->resolved_type, rhs->resolved_type);
+      CHECK(resolved.defined())
+        << "Unable to unify parent types: "
+        << lhs->resolved_type << " and " << rhs->resolved_type;
+      TypeNode* top = solver_->GetTypeNode(resolved);
+      solver_->MergeFromTo(lhs, top);
+      solver_->MergeFromTo(rhs, top);
+      return resolved;
+    }
+  }
+
+  // Checks whether lhs (taken to be a type var) occurs in t, meaning
+  // there is a recursive equality constraint, which should be rejected.
+  // N.b.: A tautology like ?a = ?a is okay and should be checked for
+  // *before* calling this method
+  bool CheckOccurs(TypeNode* lhs, const Type& t) {
+    OccursChecker rc(solver_, lhs);
+    return rc.Check(t);
+  }
+
+  // default: unify only if alpha-equal
+  Type VisitTypeDefault_(const Node* op, const Type& tn) override {
+    NodeRef nr = GetRef<NodeRef>(op);
+    Type t1 = GetRef<Type>(nr.as_derived<tvm::relay::TypeNode>());
+    if (!AlphaEqual(t1, tn)) {
+      return Type(nullptr);
+    }
+    return t1;
+  }
+
+  Type VisitType_(const TupleTypeNode* op, const Type& tn) override {
+    const auto* ttn = tn.as<TupleTypeNode>();
+    if (!ttn || op->fields.size() != ttn->fields.size()) {
+      return Type(nullptr);
+    }
+
+    TupleType tt1 = GetRef<TupleType>(op);
+    TupleType tt2 = GetRef<TupleType>(ttn);
+
+    std::vector<Type> new_fields;
+    for (size_t i = 0; i < tt1->fields.size(); i++) {
+      Type field = Unify(tt1->fields[i], tt2->fields[i]);
+      new_fields.push_back(field);
+    }
+    return TupleTypeNode::make(new_fields);
+  }
+
+  Type VisitType_(const FuncTypeNode* op, const Type& tn) override {
+    const auto* ftn = tn.as<FuncTypeNode>();
+    if (!ftn
+        || op->arg_types.size() != ftn->arg_types.size()
+        || op->type_params.size() != ftn->type_params.size()
+        || op->type_constraints.size() != ftn->type_constraints.size()) {
+      return Type(nullptr);
+    }
+
+    // remap type vars so they match
+    Map<TypeVar, Type> subst_map;
+    for (size_t i = 0; i < op->type_params.size(); i++) {
+      subst_map.Set(ftn->type_params[i], op->type_params[i]);
+    }
+
+    auto ft1 = GetRef<FuncType>(op);
+    auto ft2 = Downcast<FuncType>(Bind(GetRef<FuncType>(ftn), subst_map));
+
+    Type ret_type = Unify(ft1->ret_type, ft2->ret_type);
+
+    std::vector<Type> arg_types;
+    for (size_t i = 0; i < ft1->arg_types.size(); i++) {
+      Type arg_type = Unify(ft1->arg_types[i], ft2->arg_types[i]);
+      arg_types.push_back(arg_type);
+    }
+
+    std::vector<TypeConstraint> type_constraints;
+    for (size_t i = 0; i < ft1->type_constraints.size(); i++) {
+      Type unified_constraint = Unify(ft1->type_constraints[i],
+                                      ft2->type_constraints[i]);
+      const auto* tcn = unified_constraint.as<TypeConstraintNode>();
+      CHECK(tcn) << "Two type constraints unified into a non-constraint?"
+                 << ft1->type_constraints[i] << " and " << ft2->type_constraints[i];
+      type_constraints.push_back(GetRef<TypeConstraint>(tcn));
+    }
+
+    return FuncTypeNode::make(arg_types, ret_type, ft1->type_params, type_constraints);
+  }
+
+ private:
+  TypeSolver* solver_;
+};
+
+class TypeSolver::Resolver : public TypeMutator {
+ public:
+  explicit Resolver(TypeSolver* solver) : solver_(solver) {}
+
+  Type Resolve(const Type& t) {
+    if (!t.defined()) {
+      return t;
+    }
+    return VisitType(t);
+  }
+
+  Type VisitType_(const IncompleteTypeNode* op) override {
+    auto* node = solver_->GetTypeNode(GetRef<IncompleteType>(op));
+    return node->resolved_type;
+  }
+
+ private:
+  TypeSolver* solver_;
+};
+
+// It ends up being more compact to simply have TypeFunctor<void(const Type&) than
+// a TypeVisitor because we can use the default case to dispense with
+// most of the overrides.
+class TypeSolver::Propagator : public TypeFunctor<void(const Type&)> {
+ public:
+  explicit Propagator(TypeSolver* solver, const std::unordered_set<RelationNode*>* rels)
+    : solver_(solver), rels_(rels) {}
+
+  // adds the relation node to t and all child types of t
+  void Propagate(const Type& t) {
+    VisitType(t);
+  }
+
+  void UpdateRelSet(const Type& t) {
+    TypeNode* tnode = solver_->GetTypeNode(t);
+    for (auto* rel : *rels_) {
+      tnode->rel_set.insert(rel);
+    }
+  }
+
+  void VisitTypeDefault_(const Node* op) override {
+    NodeRef nr = GetRef<NodeRef>(op);
+    Type t = GetRef<Type>(nr.as_derived<tvm::relay::TypeNode>());
+    UpdateRelSet(t);
+  }
+
+  void VisitType_(const TupleTypeNode* op) override {
+    TupleType tt = GetRef<TupleType>(op);
+    UpdateRelSet(tt);
+
+    for (const Type& t : tt->fields) {
+      Propagate(t);
+    }
+  }
+
+  void VisitType_(const FuncTypeNode* op) override {
+    FuncType ft = GetRef<FuncType>(op);
+    UpdateRelSet(ft);
+
+    Propagate(ft->ret_type);
+    for (auto arg_type : ft->arg_types) {
+      Propagate(arg_type);
+    }
+
+    for (auto type_param : ft->type_params) {
+      Propagate(type_param);
+    }
+
+    for (auto type_cs : ft->type_constraints) {
+      Propagate(type_cs);
+    }
+  }
+
+ private:
+  TypeSolver* solver_;
+  const std::unordered_set<RelationNode*>* rels_;
+};
+
+// similarly, we use TypeFunctor<void(const Type&)> so we can use
+// the default visitor case to avoid more overrides
+class TypeSolver::Merger : public TypeFunctor<void(const Type&)> {
+ public:
+  explicit Merger(TypeSolver* solver) : solver_(solver) {}
+
+  // Merges src node to dst, ensures *all* type relations of all
+  // child nodes of src are transferred to dst.
+  void Merge(TypeNode* src, TypeNode* dst) {
+    if (src == dst) return;
+    dst_ = dst;
+    VisitType(src->resolved_type);
+    // set parent at the end so later calls to GetTypeNode go back to src
+    src->parent = dst;
+
+    // now propagate relations to child nodes, since change to
+    // a child node should update parent too
+    Propagator prop(solver_, &dst->rel_set);
+    prop.Propagate(dst->resolved_type);
+  }
+
+  // Transfers any relations linked to t to the stored dst.
+  // Any unresolved relations are added back to the queue, since
+  // there is now new information
+  void TransferLinks(const Type& t) {
+    TypeNode* src = solver_->GetTypeNode(t);
+    if (src == dst_) return;
+    for (auto* rel : src->rel_set) {
+      // if the relation is not yet resolved, add to queue
+      if (!rel->resolved) {
+        solver_->AddToQueue(rel);
+        dst_->rel_set.insert(rel);
+      }
+    }
+  }
+
+  void VisitTypeDefault_(const Node* op) override {
+    NodeRef nr = GetRef<NodeRef>(op);
+    Type t = GetRef<Type>(nr.as_derived<tvm::relay::TypeNode>());
+    TransferLinks(t);
+  }
+
+  void VisitType_(const TupleTypeNode* ttn) override {
+    auto tup = GetRef<TupleType>(ttn);
+    TransferLinks(tup);
+
+    for (auto field : tup->fields) {
+      VisitType(field);
+    }
+  }
+
+  void VisitType_(const FuncTypeNode* ftn) override {
+    auto func = GetRef<FuncType>(ftn);
+    TransferLinks(func);
+
+    VisitType(func->ret_type);
+    for (auto arg : func->arg_types) {
+      VisitType(arg);
+    }
+    for (auto param : func->type_params) {
+      VisitType(param);
+    }
+    for (auto constraint : func->type_constraints) {
+      VisitType(constraint);
+    }
+  }
+
+ private:
+  TypeSolver* solver_;
+  TypeNode* dst_;
+};
+
 // constructor
 TypeSolver::TypeSolver()
-    : reporter_(make_node<Reporter>(this)) {
+  : reporter_(make_node<Reporter>(this)) {
 }
 
 // destructor
@@ -54,31 +344,16 @@ TypeSolver::~TypeSolver() {
   }
 }
 
+// merge src type node to dst
+void TypeSolver::MergeFromTo(TypeNode* src, TypeNode* dst) {
+  Merger merger(this);
+  merger.Merge(src, dst);
+}
+
 // Add equality constraint
 Type TypeSolver::Unify(const Type& dst, const Type& src) {
-  // Known limitation
-  // - handle composite types whose component can be unknown.
-  // - handle shape pattern matching
-  TypeNode* lhs = GetTypeNode(dst);
-  TypeNode* rhs = GetTypeNode(src);
-
-  // do occur check so we don't create self-referencing structure
-  if (lhs->FindRoot() == rhs->FindRoot()) {
-    return lhs->resolved_type;
-  }
-  if (lhs->resolved_type.as<IncompleteTypeNode>()) {
-    MergeFromTo(lhs, rhs);
-    return rhs->resolved_type;
-  } else if (rhs->resolved_type.as<IncompleteTypeNode>()) {
-    MergeFromTo(rhs, lhs);
-    return lhs->resolved_type;
-  } else {
-    lhs->parent = rhs;
-    CHECK(AlphaEqual(lhs->resolved_type, rhs->resolved_type))
-        << "Incompatible parent types in UF:"
-        << lhs->resolved_type << " and " << rhs->resolved_type;
-    return rhs->resolved_type;
-  }
+  Unifier unifier(this);
+  return unifier.Unify(dst, src);
 }
 
 // Add type constraint to the solver.
@@ -96,9 +371,9 @@ void TypeSolver::AddConstraint(const TypeConstraint& constraint) {
       tlink->value = tnode;
       rnode->type_list.Push(tlink);
       // insert type->relation node
-      LinkNode<RelationNode*>* rlink = arena_.make<LinkNode<RelationNode*> >();
-      rlink->value = rnode;
-      tnode->rel_list.Push(rlink);
+      std::unordered_set<RelationNode*> singleton { rnode };
+      Propagator prop(this, &singleton);
+      prop.Propagate(tnode->resolved_type);
     }
     // add the relation to the working queue.
     this->AddToQueue(rnode);
@@ -110,12 +385,10 @@ void TypeSolver::AddConstraint(const TypeConstraint& constraint) {
 
 // Resolve a type in the solver context.
 Type TypeSolver::Resolve(const Type& type) {
+  Resolver resolver(this);
   auto it = tmap_.find(type);
-  if (it != tmap_.end()) {
-    return it->second->FindRoot()->resolved_type;
-  } else {
-    return type;
-  }
+  Type t = (it != tmap_.end()) ? it->second->FindRoot()->resolved_type : type;
+  return resolver.Resolve(t);
 }
 
 bool TypeSolver::Solve() {
@@ -128,7 +401,7 @@ bool TypeSolver::Solve() {
     // update the relation with given evidence.
     Array<Type> args;
     for (auto* tlink = rnode->type_list.head; tlink != nullptr; tlink = tlink->next) {
-      args.push_back(tlink->value->FindRoot()->resolved_type);
+      args.push_back(Resolve(tlink->value->FindRoot()->resolved_type));
       CHECK_LE(args.size(), rel->args.size());
     }
     // call the function
@@ -161,8 +434,8 @@ TVM_REGISTER_API("relay._ir_pass._test_type_solver")
             return solver->Solve();
           });
       } else if (name == "Unify") {
-        return TypedPackedFunc<void(Type, Type)>([solver](Type lhs, Type rhs) {
-            solver->Unify(lhs, rhs);
+        return TypedPackedFunc<Type(Type, Type)>([solver](Type lhs, Type rhs) {
+            return solver->Unify(lhs, rhs);
           });
       } else if (name == "Resolve") {
         return TypedPackedFunc<Type(Type)>([solver](Type t) {

src/relay/pass/type_solver.h

@@ -18,6 +18,7 @@ namespace relay {
 
 using common::LinkNode;
 using common::LinkedList;
+
 /*!
  * \brief Interface of type solver used in type inference.
  *
@@ -65,6 +66,11 @@ class TypeSolver {
   Type Unify(const Type& lhs, const Type& rhs);
 
  private:
+  class OccursChecker;
+  class Unifier;
+  class Resolver;
+  class Propagator;
+  class Merger;
   class Reporter;
   struct TypeNode;
   struct RelationNode;
@@ -77,15 +83,15 @@ class TypeSolver {
    *  that can unifies the same types to the name resolved_type.
    *
    *  It also contains collection of links to related Relations,
-   *  which is stored in rel_list.
+   *  which is stored in rel_set.
    */
   struct TypeNode {
     /*! \brief The final resolved type */
     Type resolved_type;
     /*! \brief type node in the union find algorithm */
     TypeNode* parent{nullptr};
-    /*! \brief list of relations that is related to this type node */
-    LinkedList<RelationNode*> rel_list;
+    /*! \brief set of relations that is related to this type node */
+    std::unordered_set<RelationNode*> rel_set;
     /*!
      * \brief Find the root type node, perform path compression
      * \return The root type node.
@@ -125,7 +131,7 @@ class TypeSolver {
   size_t num_resolved_rels_{0};
   /*! \brief map from type node to types. */
   std::unordered_map<Type, TypeNode*, NodeHash, NodeEqual> tmap_;
-  /*! \breif Internal queue to update the relation */
+  /*! \brief Internal queue to update the relation */
   std::queue<RelationNode*> update_queue_;
   /*! \brief allocator of all the internal node obhect*/
   common::Arena arena_;
@@ -163,22 +169,7 @@ class TypeSolver {
    * \param src The source operand
    * \param dst The dst operand.
    */
-  void MergeFromTo(TypeNode* src, TypeNode* dst) {
-    if (src == dst) return;
-    src->parent = dst;
-    // move the link to the to dst
-    for (auto* rlink = src->rel_list.head; rlink != nullptr;) {
-      // store next pointer first before rlink get moved
-      auto* next = rlink->next;
-      // if the relation is not yet resolved
-      // send the relation to the new
-      if (!rlink->value->resolved) {
-        this->AddToQueue(rlink->value);
-        dst->rel_list.Push(rlink);
-      }
-      rlink = next;
-    }
-  }
+  void MergeFromTo(TypeNode* src, TypeNode* dst);
 };
 
 }  // namespace relay

src/relay/pass/util.cc

@@ -12,105 +12,211 @@
 namespace tvm {
 namespace relay {
 
-// FreeTypeVar
-class FreeTypeVarTVisitor : public TypeVisitor {
+template<typename T>
+struct InsertionSet {
+  std::unordered_set<T, NodeHash, NodeEqual> set;
+  std::vector<T> data;
+  void Insert(const T& t) {
+    if (set.count(t) == 0) {
+      set.insert(t);
+      data.push_back(t);
+    }
+  }
+};
+
+class TypeVarTVisitor : public TypeVisitor {
  public:
-  FreeTypeVarTVisitor(
-      Array<TypeVar>* free_vars,
-      std::unordered_set<TypeVar, NodeHash, NodeEqual>* bound_vars)
-      : free_vars_(free_vars), bound_vars_(bound_vars) { }
+  TypeVarTVisitor(
+      InsertionSet<TypeVar>* type_vars,
+      InsertionSet<TypeVar>* bound_type_vars)
+      : type_vars_(type_vars), bound_type_vars_(bound_type_vars) { }
 
   void VisitType_(const TypeVarNode* tp) final {
     TypeVar var = GetRef<TypeVar>(tp);
-    if (bound_vars_->count(var) == 0) {
-      free_vars_->push_back(var);
-    }
+    type_vars_->Insert(var);
   }
 
   void VisitType_(const FuncTypeNode* f) final {
     for (auto type_param : f->type_params) {
-      bound_vars_->insert(type_param);
+      type_vars_->Insert(type_param);
+      bound_type_vars_->Insert(type_param);
     }
     TypeVisitor::VisitType_(f);
   }
 
  private:
-  Array<TypeVar>* free_vars_;
-  std::unordered_set<TypeVar, NodeHash, NodeEqual>* bound_vars_;
+  InsertionSet<TypeVar>* type_vars_;
+  InsertionSet<TypeVar>* bound_type_vars_;
 };
 
-class FreeTypeVarEVisitor : private ExprVisitor {
+class TypeVarEVisitor : private ExprVisitor {
  public:
-  Array<TypeVar> Find(const Expr& expr) {
-    this->VisitExpr(expr);
-    return free_vars_;
+  Array<TypeVar> CollectFree() {
+    Array<TypeVar> ret;
+    for (const auto& v : type_vars_.data) {
+      if (bound_type_vars_.set.count(v) == 0) {
+        ret.push_back(v);
+      }
+    }
+    return ret;
+  }
+
+  Array<TypeVar> CollectBound() {
+    Array<TypeVar> ret;
+    for (const auto& v : bound_type_vars_.data) {
+      ret.push_back(v);
+    }
+    return ret;
+  }
+
+  Array<TypeVar> CollectAll() {
+    Array<TypeVar> ret;
+    for (const auto& v : type_vars_.data) {
+      ret.push_back(v);
+    }
+    return ret;
   }
 
-  Array<TypeVar> Find(const Type& type) {
-    this->VisitType(type);
-    return free_vars_;
+  Array<TypeVar> Free(const Expr& expr) {
+    VisitExpr(expr);
+    return CollectFree();
+  }
+
+  Array<TypeVar> Free(const Type& type) {
+    VisitType(type);
+    return CollectFree();
+  }
+
+  Array<TypeVar> Bound(const Expr& expr) {
+    VisitExpr(expr);
+    return CollectBound();
+  }
+
+  Array<TypeVar> Bound(const Type& type) {
+    VisitType(type);
+    return CollectBound();
+  }
+
+  Array<TypeVar> All(const Expr& expr) {
+    VisitExpr(expr);
+    return CollectAll();
+  }
+
+  Array<TypeVar> All(const Type& type) {
+    VisitType(type);
+    return CollectAll();
   }
 
   void VisitExpr_(const FunctionNode* f) final {
     for (const auto& tp : f->type_params) {
-      bound_vars_.insert(tp);
+      type_vars_.Insert(tp);
+      bound_type_vars_.Insert(tp);
     }
     ExprVisitor::VisitExpr_(f);
   }
 
   void VisitType(const Type& t) final {
-    FreeTypeVarTVisitor(&free_vars_, &bound_vars_)
+    TypeVarTVisitor(&type_vars_, &bound_type_vars_)
         .VisitType(t);
   }
 
  private:
-  // The result list
-  Array<TypeVar> free_vars_;
-  std::unordered_set<TypeVar, NodeHash, NodeEqual> bound_vars_;
+  InsertionSet<TypeVar> type_vars_;
+  InsertionSet<TypeVar> bound_type_vars_;
 };
 
-class FreeVarVisitor : protected ExprVisitor {
+class VarVisitor : protected ExprVisitor {
  public:
-  Array<Var> Find(const Expr& expr) {
+  Array<Var> Free(const Expr& expr) {
     this->VisitExpr(expr);
-    return free_vars_;
+    Array<Var> ret;
+    for (const auto& v : vars_.data) {
+      if (bound_vars_.set.count(v) == 0) {
+        ret.push_back(v);
+      }
+    }
+    return ret;
   }
 
-  void VisitExpr_(const VarNode* var) final {
-    if (bound_vars_.count(var) == 0) {
-      free_vars_.push_back(GetRef<Var>(var));
+  Array<Var> Bound(const Expr& expr) {
+    this->VisitExpr(expr);
+    Array<Var> ret;
+    for (const auto& v : bound_vars_.data) {
+      ret.push_back(v);
     }
+    return ret;
+  }
+
+  Array<Var> All(const Expr& expr) {
+    this->VisitExpr(expr);
+    Array<Var> ret;
+    for (const auto& v : vars_.data) {
+      ret.push_back(v);
+    }
+    return ret;
+  }
+
+  void MarkBounded(const Var& v) {
+    bound_vars_.Insert(v);
+    vars_.Insert(v);
+  }
+
+  void VisitExpr_(const VarNode* var) final {
+    vars_.Insert(GetRef<Var>(var));
   }
 
   void VisitExpr_(const FunctionNode* op) final {
     for (const auto& param : op->params) {
-      bound_vars_.insert(param.operator->());
+      MarkBounded(param);
     }
     VisitExpr(op->body);
   }
 
   void VisitExpr_(const LetNode* op) final {
-    bound_vars_.insert(op->var.operator->());
+    MarkBounded(op->var);
     VisitExpr(op->value);
     VisitExpr(op->body);
   }
 
  private:
-  // The result list
-  Array<Var> free_vars_;
-  std::unordered_set<const VarNode*> bound_vars_;
+  InsertionSet<Var> vars_;
+  InsertionSet<Var> bound_vars_;
 };
 
 tvm::Array<TypeVar> FreeTypeVars(const Expr& expr) {
-  return FreeTypeVarEVisitor().Find(expr);
+  return TypeVarEVisitor().Free(expr);
 }
 
 tvm::Array<TypeVar> FreeTypeVars(const Type& type) {
-  return FreeTypeVarEVisitor().Find(type);
+  return TypeVarEVisitor().Free(type);
+}
+
+tvm::Array<TypeVar> BoundTypeVars(const Expr& expr) {
+  return TypeVarEVisitor().Bound(expr);
+}
+
+tvm::Array<TypeVar> BoundTypeVars(const Type& type) {
+  return TypeVarEVisitor().Bound(type);
+}
+
+tvm::Array<TypeVar> AllTypeVars(const Expr& expr) {
+  return TypeVarEVisitor().All(expr);
+}
+
+tvm::Array<TypeVar> AllTypeVars(const Type& type) {
+  return TypeVarEVisitor().All(type);
 }
 
 tvm::Array<Var> FreeVars(const Expr& expr) {
-  return FreeVarVisitor().Find(expr);
+  return VarVisitor().Free(expr);
+}
+
+tvm::Array<Var> BoundVars(const Expr& expr) {
+  return VarVisitor().Bound(expr);
+}
+
+tvm::Array<Var> AllVars(const Expr& expr) {
+  return VarVisitor().All(expr);
 }
 
 TVM_REGISTER_API("relay._ir_pass.free_vars")
@@ -118,16 +224,46 @@ TVM_REGISTER_API("relay._ir_pass.free_vars")
     *ret = FreeVars(args[0]);
   });
 
+TVM_REGISTER_API("relay._ir_pass.bound_vars")
+  .set_body([](TVMArgs args, TVMRetValue* ret) {
+      *ret = BoundVars(args[0]);
+    });
+
+TVM_REGISTER_API("relay._ir_pass.all_vars")
+  .set_body([](TVMArgs args, TVMRetValue* ret) {
+      *ret = AllVars(args[0]);
+    });
+
 TVM_REGISTER_API("relay._ir_pass.free_type_vars")
 .set_body([](TVMArgs args, TVMRetValue* ret) {
     NodeRef x = args[0];
-    if (x.as<TypeNode>()) {
+    if (x.as_derived<TypeNode>()) {
       *ret = FreeTypeVars(Downcast<Type>(x));
     } else {
       *ret = FreeTypeVars(Downcast<Expr>(x));
     }
   });
 
+TVM_REGISTER_API("relay._ir_pass.bound_type_vars")
+  .set_body([](TVMArgs args, TVMRetValue* ret) {
+      NodeRef x = args[0];
+      if (x.as_derived<TypeNode>()) {
+        *ret = BoundTypeVars(Downcast<Type>(x));
+      } else {
+        *ret = BoundTypeVars(Downcast<Expr>(x));
+      }
+    });
+
+TVM_REGISTER_API("relay._ir_pass.all_type_vars")
+  .set_body([](TVMArgs args, TVMRetValue* ret) {
+      NodeRef x = args[0];
+      if (x.as_derived<TypeNode>()) {
+        *ret = AllTypeVars(Downcast<Type>(x));
+      } else {
+        *ret = AllTypeVars(Downcast<Expr>(x));
+      }
+    });
+
 /*!
  * \brief Get reference counter of each internal ExprNode in body.
  * \param body The body expression.

tests/cpp/relay_pass_type_infer_test.cc

@@ -6,13 +6,17 @@
 
 TEST(Relay, SelfReference) {
   using namespace tvm;
-  auto type_a = relay::TypeVarNode::make("a", relay::TypeVarNode::kType);
-  auto type_b = relay::TypeVarNode::make("b", relay::TypeVarNode::kType);
-  auto x = relay::VarNode::make("x", type_a);
-  auto f = relay::FunctionNode::make(tvm::Array<relay::Var>{ x }, x, type_b, Array<relay::TypeVar>{});
-  auto fx = relay::CallNode::make(f, Array<relay::Expr>{ x });
+  auto tensor_type = relay::TensorTypeNode::make({}, ::tvm::Bool());
+  auto x = relay::VarNode::make("x", relay::Type());
+  auto f = relay::FunctionNode::make(tvm::Array<relay::Var>{ x }, x, relay::Type(), {});
+
+  auto y = relay::VarNode::make("y", tensor_type);
+  auto call = relay::CallNode::make(f, Array<relay::Expr>{ y });
+  auto fx = relay::FunctionNode::make(tvm::Array<relay::Var>{ y }, call, relay::Type(), {});
   auto type_fx = relay::InferType(fx, relay::ModuleNode::make(Map<relay::GlobalVar, relay::Function>{}));
-  CHECK_EQ(type_fx->checked_type(), type_a);
+
+  auto expected = relay::FuncTypeNode::make(tvm::Array<relay::Type>{ tensor_type }, tensor_type, {}, {});
+  CHECK(AlphaEqual(type_fx->checked_type(), expected));
 }
 
 int main(int argc, char ** argv) {

tests/python/relay/test_pass_free_vars.py

-import tvm
-from tvm import relay
-from tvm.relay.ir_pass import free_vars, free_type_vars
-
-def test_free_vars():
-    ty = relay.TensorType([], "int32")
-    x = relay.Var("x", ty)
-    fvx = free_vars(x)
-    assert len(fvx) == 1
-    assert fvx[0] == x
-    v = relay.Constant(tvm.nd.array(10))
-
-    let = relay.Let(x, v, x)
-    fvx = free_vars(let)
-    assert len(free_vars(let)) == 0
-    f = relay.Function([x], x, ty)
-    assert len(free_vars(f)) == 0
-
-
-def test_tuple():
-    t = relay.Var('t')
-    fv = free_vars(relay.Tuple([t, t]))
-    assert len(fv) == 1
-    assert fv[0] == t
-    fv = free_vars(relay.TupleGetItem(t, 123))
-    assert len(fv) == 1
-    assert fv[0] == t
-
-
-def test_free_type_vars():
-    tp = relay.TypeVar("")
-    ty = relay.TupleType([tp, relay.TensorType([], "int32")])
-    x = relay.Var("x", ty)
-    y = relay.Var("y")
-    let = relay.Let(x, y, x)
-    fvl = free_vars(let)
-    assert len(fvl) == 1
-    assert fvl[0] == y
-    ftvl = free_type_vars(let)
-    assert len(ftvl) == 1
-    assert ftvl[0] == tp

tests/python/relay/test_pass_vars.py

+import tvm
+from tvm import relay
+from tvm.relay.ir_pass import (free_vars, free_type_vars,
+                               bound_vars, bound_type_vars,
+                               all_vars, all_type_vars)
+
+def assert_vars_match(actual, expected):
+    assert len(actual) == len(expected)
+    for i in range(len(actual)):
+        assert actual[i] == expected[i]
+
+
+def test_free_vars():
+    ty = relay.TensorType([], "int32")
+    x = relay.Var("x", ty)
+    fvx = free_vars(x)
+    assert len(fvx) == 1
+    assert fvx[0] == x
+    v = relay.Constant(tvm.nd.array(10))
+
+    let = relay.Let(x, v, x)
+    fvx = free_vars(let)
+    assert len(free_vars(let)) == 0
+    f = relay.Function([x], x, ty)
+    assert len(free_vars(f)) == 0
+
+
+def test_free_vars_tuple():
+    t = relay.Var('t')
+    fv = free_vars(relay.Tuple([t, t]))
+    assert len(fv) == 1
+    assert fv[0] == t
+    fv = free_vars(relay.TupleGetItem(t, 123))
+    assert len(fv) == 1
+    assert fv[0] == t
+
+
+def test_free_type_vars():
+    tp = relay.TypeVar("")
+    ty = relay.TupleType([tp, relay.TensorType([], "int32")])
+    x = relay.Var("x", ty)
+    y = relay.Var("y")
+    let = relay.Let(x, y, x)
+    fvl = free_vars(let)
+    assert len(fvl) == 1
+    assert fvl[0] == y
+    ftvl = free_type_vars(let)
+    assert len(ftvl) == 1
+    assert ftvl[0] == tp
+
+
+def test_bound_vars():
+    x = relay.Var("x")
+    y = relay.Var("y")
+    z = relay.Var("z")
+    a = relay.Var("a")
+
+    f1 = relay.Function([x, y, z], relay.Let(a, x, relay.Tuple([])))
+    assert_vars_match(bound_vars(f1), [x, y, z, a])
+
+    tup = relay.Tuple([x, y, z, a])
+    assert len(bound_vars(tup)) == 0
+
+    f2 = relay.Function([x, y], relay.Tuple([x, y, z, a]))
+    assert_vars_match(bound_vars(f2), [x, y])
+
+
+def test_bound_type_vars():
+    a = relay.TypeVar("a")
+    b = relay.TypeVar("b")
+    c = relay.TypeVar("c")
+
+    ft1 = relay.FuncType([a], b, [a, b])
+    bound_ft1 = bound_type_vars(ft1)
+    assert_vars_match(bound_type_vars(ft1), [a, b])
+
+    ft2 = relay.FuncType([], c, [a])
+    assert_vars_match(bound_type_vars(ft2), [a])
+
+    tup_ty = relay.TupleType([a, b, c])
+    assert len(bound_type_vars(tup_ty)) == 0
+
+    f1 = relay.Function([], relay.Tuple([]), type_params=[a, b])
+    assert_vars_match(bound_type_vars(f1), [a, b])
+
+    f2 = relay.Function([], relay.Tuple([]), c)
+    assert len(bound_type_vars(f2)) == 0
+
+    x = relay.Var("x", a)
+    let1 = relay.Let(x, relay.Tuple([]), x)
+    assert len(bound_type_vars(let1)) == 0
+
+    let2 = relay.Let(x, relay.Function([], relay.Tuple([]), type_params=[b, c]), x)
+    assert_vars_match(bound_type_vars(let2), [b, c])
+
+
+def test_all_vars():
+    x = relay.Var("x")
+    y = relay.Var("y")
+    z = relay.Var("z")
+
+    f1 = relay.Function([x, y], z)
+    assert_vars_match(all_vars(f1), [x, y, z])
+
+    f2 = relay.Function([x], relay.Let(y, relay.Tuple([]), z))
+    assert_vars_match(all_vars(f2), [x, y, z])
+
+    f3 = relay.Function([x], relay.Tuple([y, z]))
+    assert_vars_match(all_vars(f3), [x, y, z])
+
+    tup = relay.Tuple([x, y, z])
+    assert_vars_match(all_vars(tup), [x, y, z])
+
+
+def test_all_type_vars():
+    a = relay.TypeVar("a")
+    b = relay.TypeVar("b")
+    c = relay.TypeVar("c")
+
+    ft1 = relay.FuncType([b], c, [a])
+    assert_vars_match(all_type_vars(ft1), [a, b, c])
+
+    ft2 = relay.FuncType([], relay.TupleType([a, b, c]), [])
+    assert_vars_match(all_type_vars(ft2), [a, b, c])
+
+    w = relay.Var("w")
+    x = relay.Var("x", a)
+    y = relay.Var("y", b)
+    z = relay.Var("z", c)
+    
+    f1 = relay.Function([x], y, b, [a])
+    assert_vars_match(all_type_vars(f1), [a, b])
+
+    f2 = relay.Function([x], relay.Let(y, x, z))
+    assert_vars_match(all_type_vars(f2), [a, b, c])
+
+    f3 = relay.Function([], relay.Tuple([x, y, z]), ret_type=relay.TupleType([a, b, c]))
+    assert_vars_match(all_type_vars(f3), [a, b, c])
+
+    f4 = relay.Function([w], relay.Tuple([]), type_params=[a, b, c])
+    assert_vars_match(all_type_vars(f4), [a, b, c])
+
+    f5 = relay.Function([w], w)
+    assert len(all_type_vars(f5)) == 0

tests/python/relay/test_type_infer.py

@@ -23,7 +23,7 @@ def test_monomorphic_let():
     x = sb.let('x', relay.const(1.0, "float64"))
     sb.ret(x)
     xchecked = relay.ir_pass.infer_type(sb.get())
-    assert xchecked.checked_type == relay.scalar_type("float64")
+    assert xchecked.checked_type == relay.scalar_type("float64" )
 
 
 def test_single_op():
@@ -41,14 +41,15 @@ def test_add_broadcast_op():
             return x + y;
         }
     """
-    pass
-    # x = relay.var('x', shape=(10, 4))
-    # y = relay.var('y', shape=(5, 10, 1))
-    # z = x + y
-    # func = relay.Function([x, y], z)
-    # ttype = relay.TensorType((5, 5, 5), 'float32')
-    # expected_ty = relay.FuncType([ttype, ttype], ttype)
-    # assert_has_type(func.to_func(), expected_ty)
+    x = relay.var('x', shape=(10, 4))
+    y = relay.var('y', shape=(5, 10, 1))
+    z = x + y
+    func = relay.Function([x, y], z)
+    t1 = relay.TensorType((10, 4), 'float32')
+    t2 = relay.TensorType((5, 10, 1), 'float32')
+    t3 = relay.TensorType((5, 10, 4), 'float32')
+    expected_ty = relay.FuncType([t1, t2], t3)
+    assert_has_type(func, expected_ty)
 
 
 def test_dual_op():
@@ -110,24 +111,17 @@ def test_recursion():
     assert "%3 = @f(%1, %2)" in mod.astext()
     assert mod[f].checked_type == relay.FuncType([ti32, tf32], tf32)
 
-# This currently fails and should pass under the type system.
-#
-# This test is to illustrate problem with our weak form of
-# unification.
-#
-
 
 def test_incomplete_call():
-    sb = ScopeBuilder()
-    x = relay.var('x', dtype='int32')
+    tt = relay.scalar_type('int32')
+    x = relay.var('x', tt)
     f = relay.var('f')
-    func = relay.Function([x, f], relay.Call(f, [x]))
+    func = relay.Function([x, f], relay.Call(f, [x]), tt)
+
+    ft = relay.ir_pass.infer_type(func)
+    f_type = relay.FuncType([tt], tt)
+    assert ft.checked_type == relay.FuncType([tt, f_type], tt)
 
-    try:
-        relay.ir_pass.infer_type(func)
-        assert False
-    except tvm.TVMError as e:
-        assert True
 
 def test_tuple():
     tp = relay.TensorType((10,))
@@ -136,6 +130,7 @@ def test_tuple():
     assert (relay.ir_pass.infer_type(res).checked_type ==
             relay.TupleType([tp, tp]))
 
+
 def test_free_expr():
     x = relay.var("x", "float32")
     y = relay.add(x, x)
@@ -161,38 +156,26 @@ def test_type_args():
     assert sh2[1].value == 10
 
 
-def test_self_reference():
-    """
-    Program:
-       def f(x) {
-           return x;
-       }
-    """
-    a = relay.TypeVar("a")
-    x = relay.var("x", a)
-    sb = relay.ScopeBuilder()
-
-    f = relay.Function([x], x)
-    fx = relay.Call(f, [x])
-    assert relay.ir_pass.infer_type(x).checked_type == a
-    assert relay.ir_pass.infer_type(f).checked_type == relay.FuncType([a], a)
-    assert relay.ir_pass.infer_type(fx).checked_type == a
-
-
-def test_global_var_cow_issue():
+def test_global_var_recursion():
     mod = relay.Module({})
     gv = relay.GlobalVar("foo")
     x = relay.var('x', shape=[])
-    func = relay.Function([x], relay.Call(gv, [x]),
-                          relay.TensorType([], 'float32'))
+    tt = relay.scalar_type('float32')
+
+    func = relay.Function([x], relay.Call(gv, [x]), tt)
     mod[gv] = func
 
+    ft = relay.ir_pass.infer_type(gv, mod)
+    assert mod[ft].checked_type == relay.FuncType([tt], tt)
+
 
 def test_equal():
     i = relay.var('i', shape=[], dtype='int32')
     eq = op.equal(i, relay.const(0, dtype='int32'))
-    # This should fail ....
-    func = relay.Function([i], eq, ret_type=relay.TensorType([], 'int32'))
+    func = relay.Function([i], eq)
+    ft = relay.ir_pass.infer_type(func)
+
+    assert ft.checked_type == relay.FuncType([relay.scalar_type('int32')], relay.scalar_type('bool'))
 
 
 if __name__ == "__main__":
@@ -204,8 +187,12 @@ if __name__ == "__main__":
     test_decl()
     test_recursion()
     test_tuple()
+    test_generalized_tuple()
     test_incomplete_call()
+    test_generalized_call()
+    test_call_with_type_args()
     test_free_expr()
     test_type_args()
     test_self_reference()
-    test_global_var_cow_issue()
+    test_global_var_recursion()
+    test_equal()

tests/python/relay/test_type_solver.py

 import tvm
 from tvm import relay
+from nose.tools import raises
 
 
 def make_rel(name, args, num_inputs=None, attrs=None):
@@ -48,7 +49,170 @@ def test_backward_solving():
     assert solver.Resolve(t3) == relay.ty.TensorType((10, 10, 20), "float32")
 
 
+def test_unify_tuple():
+    solver = make_solver()
+    t1 = relay.ty.IncompleteType()
+    t2 = relay.ty.IncompleteType()
+    t3 = relay.ty.TensorType((10, 20), "float32")
+
+    tup1 = relay.ty.TupleType([t1, t2])
+    tup2 = relay.ty.TupleType([t3, t3])
+
+    unified = solver.Unify(tup1, tup2)
+    assert unified == tup2
+
+
+def test_unify_functype():
+    solver = make_solver()
+    t1 = relay.ty.IncompleteType()
+    t2 = relay.ty.IncompleteType()
+    t3 = relay.ty.IncompleteType()
+
+    unit = relay.ty.TupleType([])
+    tensor1 = relay.ty.TensorType((10, 20), "float32")
+    tensor2 = relay.ty.TensorType((10,), "float32")
+
+    ft1 = relay.ty.FuncType([t1, t2], t3)
+    ft2 = relay.ty.FuncType([tensor1, tensor2], unit)
+
+    unified = solver.Unify(ft1, ft2)
+    assert unified == ft2
+
+
+def test_recursive_unify():
+    solver = make_solver()
+    t1 = relay.ty.IncompleteType()
+    t2 = relay.ty.IncompleteType()
+    t3 = relay.ty.IncompleteType()
+
+    tensor1 = relay.ty.TensorType((10, 10, 20), "float32")
+    tensor2 = relay.ty.TensorType((10, 20), "float32")
+    tensor3 = relay.ty.TensorType((10,), "float32")
+
+    tup1 = relay.ty.TupleType([relay.ty.TupleType([t1, t2]), t2])
+    tup2 = relay.ty.TupleType([relay.ty.TupleType([tensor1, tensor2]), tensor2])
+
+    ft1 = relay.ty.FuncType([tup1, t3], t3)
+    ft2 = relay.ty.FuncType([tup2, tensor3], tensor3)
+
+    unified = solver.Unify(ft1, ft2)
+    assert unified == ft2
+
+
+def test_unify_vars_under_tuples():
+    solver = make_solver()
+    t1 = relay.ty.IncompleteType()
+
+    tup1 = relay.ty.TupleType([t1, t1])
+    unified = solver.Unify(tup1, tup1)
+    assert unified == tup1
+
+    t2 = relay.ty.IncompleteType()
+    tup2 = relay.ty.TupleType([t2, t2])
+
+    tup3 = relay.ty.TupleType([t1, t2])
+    tup4 = relay.ty.TupleType([t2, t1])
+    unified = solver.Unify(tup3, tup4)
+    assert (unified == tup1 or unified == tup2)
+
+
+def test_binding_over_typevars():
+    solver = make_solver()
+
+    t1 = relay.ty.IncompleteType()
+    t2 = relay.ty.IncompleteType()
+
+    a = relay.ty.TypeVar('a')
+    b = relay.ty.TypeVar('b')
+    c = relay.ty.TypeVar('c')
+    d = relay.ty.TypeVar('d')
+
+    ft1 = relay.ty.FuncType([t1], t2, [c, d])
+    ft2 = relay.ty.FuncType([a], b, [a, b])
+    unified = solver.Unify(ft1, ft2)
+    assert (unified == solver.Resolve(ft1))
+
+
+def test_recursive_backward_solving():
+    solver = make_solver()
+
+    tensor1 = relay.ty.TensorType((10, 20), "float32")
+    tensor2 = relay.ty.TensorType((10, 1, 1), "float32")
+    tensor3 = relay.ty.TensorType((10,), "float32")
+
+    t1 = relay.ty.IncompleteType()
+    t2 = relay.ty.IncompleteType()
+    t3 = relay.ty.IncompleteType()
+
+    tup1 = relay.ty.TupleType([relay.ty.TupleType([tensor1, tensor2]), tensor3])
+    tup2 = relay.ty.TupleType([relay.ty.TupleType([t1, t2]), t3])
+    solver.gen_type("Identity", [tup1], out=tup2)
+
+    assert solver.Solve()
+    assert solver.Resolve(tup2) == tup1
+
+
+def test_backward_solving_after_child_update():
+    solver = make_solver()
+
+    tensor1 = relay.ty.TensorType((10, 20), "float32")
+    tensor2 = relay.ty.TensorType((10, 1, 1), "float32")
+
+    t1 = relay.ty.IncompleteType()
+    t2 = relay.ty.IncompleteType()
+    t3 = relay.ty.IncompleteType()
+
+    tup1 = relay.ty.TupleType([t1, t2])
+    tup2 = relay.ty.TupleType([t1, t3])
+
+    tup_concrete = relay.ty.TupleType([tensor1, tensor2])
+
+    t4 = solver.gen_type("Identity", [tup1])
+    t5 = solver.gen_type("Identity", [tup2])
+
+    solver.gen_type("Identity", [t4], out=t5)
+    assert solver.Solve()
+    assert solver.Resolve(t3) == t3 or solver.Resolve(t3) == t2
+    assert solver.Resolve(t4) == tup1 or solver.Resolve(t4) == tup2
+    assert solver.Resolve(t5) == tup1 or solver.Resolve(t5) == tup2
+
+    # updating the variables *inside* tup1 and tup2 should update t4 and t5
+    solver.gen_type("Identity", [t1], out=tensor1)
+    solver.gen_type("Identity", [t2], out=tensor2)
+    assert solver.Solve()
+    assert solver.Resolve(t4) == tup_concrete
+    assert solver.Resolve(t5) == tup_concrete
+
+@raises(tvm._ffi.base.TVMError)
+def test_incompatible_tuple_unification():
+    solver = make_solver()
+    t1 = relay.ty.IncompleteType()
+    t2 = relay.ty.IncompleteType()
+
+    tensor1 = relay.ty.TensorType((1, 2, 3), "float32")
+    tensor2 = relay.ty.TensorType((2, 3), "float32")
+    tensor3 = relay.ty.TensorType((3,), "float32")
+
+    tup1 = relay.ty.TupleType([relay.ty.TupleType([t1, t1]), t2])
+    tup2 = relay.ty.TupleType([relay.ty.TupleType([tensor1, tensor2]), tensor3])
+    solver.Unify(tup1, tup2)
+
+
+@raises(tvm._ffi.base.TVMError)
+def test_bad_recursive_unification():
+    solver = make_solver()
+    t1 = relay.ty.IncompleteType()
+    solver.Unify(t1, relay.ty.TupleType([t1, t1]))
+
 
 if __name__ == "__main__":
     test_bcast()
     test_backward_solving()
+    test_unify_tuple()
+    test_unify_functype()
+    test_recursive_unify()
+    test_unify_vars_under_tuples()
+    test_recursive_backward_solving()
+    test_backward_solving_after_child_update()
+    test_incompatible_tuple_unification()
+    test_bad_recursive_unification()