Fix old type inference invariant

src/relay/ir/module.cc

@@ -23,6 +23,7 @@ Module ModuleNode::make(tvm::Map<GlobalVar, Function> global_funcs) {
         << "Duplicate global function name " << kv.first->name_hint;
     n->global_var_map_.Set(kv.first->name_hint, kv.first);
   }
+
   n->entry_func = GlobalVarNode::make("main");
   return Module(n);
 }
@@ -104,7 +105,6 @@ Expr ModuleNode::EntryPoint() {
 Module ModuleNode::FromExpr(
   const Expr& expr,
   const tvm::Map<GlobalVar, Function>& global_funcs) {
-  GlobalVar main = GlobalVarNode::make("main");
   auto mod = ModuleNode::make({});
   auto func_node = expr.as<FunctionNode>();
   Function func;
@@ -113,7 +113,7 @@ Module ModuleNode::FromExpr(
   } else {
     func = FunctionNode::make({}, expr, Type(), {}, {});
   }
-  mod->Add(main, func);
+  mod->Add(mod->entry_func, func);
   return mod;
 }
 

src/relay/pass/type_infer.cc

@@ -603,7 +603,16 @@ Expr InferType(const Expr& expr, const Module& mod_ref) {
     // NB(@jroesch): By adding the expression to the module we will
     // type check it anyway; afterwards we can just recover type
     // from the type-checked function to avoid doing unnecessary work.
-    return mod->Lookup(mod->entry_func);
+
+    Function e = mod->Lookup(mod->entry_func);
+
+    // FromExpr wraps a naked expression as a function, we will unbox
+    // it here.
+    if (auto func = expr.as<FunctionNode>()) {
+      return e;
+    } else {
+      return e->body;
+    }
   } else {
     auto e = TypeInferencer(mod_ref, mod_ref->entry_func).Infer(expr);
     CHECK(WellFormed(e));