From f6747be9892fcdb6d62f42802446f104e366d7b9 Mon Sep 17 00:00:00 2001
From: Andrew Luo <andrew.zhao.luo@gmail.com>
Date: Mon, 13 Dec 2021 12:07:31 -0800
Subject: [PATCH 01/19] reuse checked types

---
 src/relay/op/nn/nn.h                       |  1 +
 src/relay/transforms/to_mixed_precision.cc | 18 ++++++++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/src/relay/op/nn/nn.h b/src/relay/op/nn/nn.h
index 6bc21473af18..900c61eb8200 100644
--- a/src/relay/op/nn/nn.h
+++ b/src/relay/op/nn/nn.h
@@ -162,6 +162,7 @@ bool BatchMatmulRel(const Array<Type>& types, int num_inputs, const Attrs& attrs
   if (out_dtype.bits() == 0) {
     out_dtype = x->dtype;
   }
+
   // assign output type
   const auto& out_b =
       xb->IsInstance<tir::AnyNode>() || yb->IsInstance<tir::AnyNode>() ? tir::Any() : max(xb, yb);
diff --git a/src/relay/transforms/to_mixed_precision.cc b/src/relay/transforms/to_mixed_precision.cc
index ae10c937ff1c..da0ece2ad6af 100644
--- a/src/relay/transforms/to_mixed_precision.cc
+++ b/src/relay/transforms/to_mixed_precision.cc
@@ -176,6 +176,12 @@ class MixedPrecisionPass : public MixedModeMutator {
   }
 
   Type GetType(const Expr& expr) const {
+    const Type old_checked_type = expr->checked_type_;
+
+    if (old_checked_type.defined()) {
+      return old_checked_type;
+    }
+
     auto mod = IRModule::FromExpr(expr);
     mod = transform::InferType()(mod);
     if (expr.as<FunctionNode>()) {
@@ -381,6 +387,18 @@ class MixedPrecisionPass : public MixedModeMutator {
     return Call(cur_op, new_args, pre_call_node->attrs, new_arg_types, pre_call_node->span);
   }
 
+  Expr Rewrite_(const TupleGetItemNode* pre, const Expr& post) {
+    // The old checked type in the expression may not be valid so clear it
+    post->checked_type_ = Type(nullptr);
+    return post;
+  }
+
+  Expr Rewrite_(const TupleNode* pre, const Expr& post) {
+    // The old checked type in the expression may not be valid so clear it
+    post->checked_type_ = Type(nullptr);
+    return post;
+  }
+
   Expr VisitExpr_(const FunctionNode* func) final {
     // Erase the ret_type annotation and let the normal pass recalculate
     const_cast<FunctionNode*>(func)->ret_type = Type(nullptr);

From 5d3932ff34c05e92a8a21ff55edb3a8114d144f6 Mon Sep 17 00:00:00 2001
From: Andrew Luo <andrew.zhao.luo@gmail.com>
Date: Mon, 13 Dec 2021 14:39:47 -0800
Subject: [PATCH 02/19] analogous subgraph

---
 src/relay/transforms/to_mixed_precision.cc | 30 +++++++++++++++++++---
 1 file changed, 26 insertions(+), 4 deletions(-)

diff --git a/src/relay/transforms/to_mixed_precision.cc b/src/relay/transforms/to_mixed_precision.cc
index da0ece2ad6af..de877337c2b6 100644
--- a/src/relay/transforms/to_mixed_precision.cc
+++ b/src/relay/transforms/to_mixed_precision.cc
@@ -101,6 +101,9 @@ class MixedPrecisionPass : public MixedModeMutator {
   /*! \brief The target datatype we want to convert to e.g. FP16 */
   const DataType mixed_precision_type_;
 
+  /* TODO*/
+  std::unordered_map<Expr, Expr, ObjectPtrHash, ObjectPtrEqual> analgous_graphs;
+
   /*! \brief Map of Ops with no associated FTVMMixedPrecisionConversionType to the times they were
    * encountered. Used for emitting warnings on missing ops in the pass.
    */
@@ -175,20 +178,39 @@ class MixedPrecisionPass : public MixedModeMutator {
     return Attrs(new_attrs);
   }
 
+  Expr MakeAnalogousSubgraph(const Expr& expr) const {
+    if (auto node = expr.as<CallNode>()) {
+      Array<Expr> args;
+      for (Expr expr : node->args) {
+        args.push_back(Var("dummy_temp", GetType(expr)));
+      }
+      return Call(node->op, args, node->attrs, node->type_args, node->span);
+    } else if (auto node = expr.as<TupleGetItemNode>()) {
+      return TupleGetItem(MakeAnalogousSubgraph(node->tuple), node->index, node->span);
+    } else {
+      LOG(FATAL) << "Unknown node " << expr;
+      return Expr(nullptr);
+    }
+  }
+
   Type GetType(const Expr& expr) const {
     const Type old_checked_type = expr->checked_type_;
-
     if (old_checked_type.defined()) {
       return old_checked_type;
     }
 
-    auto mod = IRModule::FromExpr(expr);
+    auto mod = IRModule::FromExpr(MakeAnalogousSubgraph(expr));
+    // LOG(WARNING) << mod;
+    // LOG(WARNING) << IRModule::FromExpr(expr);
     mod = transform::InferType()(mod);
+    Type t;
     if (expr.as<FunctionNode>()) {
-      return mod->Lookup("main")->checked_type();
+      t = mod->Lookup("main")->checked_type();
     } else {
-      return mod->Lookup("main").as<FunctionNode>()->body->checked_type();
+      t = mod->Lookup("main").as<FunctionNode>()->body->checked_type();
     }
+    expr->checked_type_ = t;
+    return t;
   }
 
   bool IsMixedPrecisionType(const Type& t, bool ignore_non_float = false) const {

From 3220c80959559ab1bcdcaea33ca8eab3caeb5e07 Mon Sep 17 00:00:00 2001
From: Andrew Luo <andrew.zhao.luo@gmail.com>
Date: Mon, 13 Dec 2021 16:58:47 -0800
Subject: [PATCH 03/19] brr go fast

---
 include/tvm/relay/transform.h              |  5 ++
 src/relay/transforms/to_mixed_precision.cc | 89 ++++++++++++++++++----
 src/relay/transforms/type_infer.cc         | 82 ++++++++++++++++++++
 3 files changed, 160 insertions(+), 16 deletions(-)

diff --git a/include/tvm/relay/transform.h b/include/tvm/relay/transform.h
index 2d6cdeaa8ca1..a0624a890858 100644
--- a/include/tvm/relay/transform.h
+++ b/include/tvm/relay/transform.h
@@ -258,6 +258,11 @@ TVM_DLL Pass DynamicToStatic();
  */
 TVM_DLL Pass InferType();
 
+/*
+TODO
+*/
+TVM_DLL Type InferTypeFast(const Expr& expr);
+
 /*!
  * \brief Search and eliminate common subexpression. For example, if there are
  * two expressions evaluated to an identical value, a single variable is created
diff --git a/src/relay/transforms/to_mixed_precision.cc b/src/relay/transforms/to_mixed_precision.cc
index de877337c2b6..ef8a78c330aa 100644
--- a/src/relay/transforms/to_mixed_precision.cc
+++ b/src/relay/transforms/to_mixed_precision.cc
@@ -36,6 +36,73 @@
 namespace tvm {
 namespace relay {
 
+class SameTypedSubgraphExtractor : public ExprMutator {
+  /*
+  Creates a small subgraph with the same type as the input expression.
+
+  ExprMutator is sufficient over MixedModemutator since we will not recurse much.
+  */
+
+  Expr VisitExpr_(const VarNode* op) { return Var(op->vid, op->type_annotation, op->span); }
+  Expr VisitExpr_(const ConstantNode* op) { return Constant(op->data, op->span); }
+  Expr VisitExpr_(const GlobalVarNode* op) { return GlobalVar(op->name_hint); }
+  Expr VisitExpr_(const OpNode* op) { return Op(GetRef<Op>(op)); }
+  Expr VisitExpr_(const TupleNode* op) {
+    return Tuple(get_analogous_expression(op->fields), op->span);
+  }
+  Expr VisitExpr_(const FunctionNode* op) {
+    // Here will be the only VisitExpr
+    return Function(op->params, get_analogous_expression(op->body), op->ret_type, op->type_params,
+                    op->attrs, op->span);
+  }
+  Expr VisitExpr_(const CallNode* op) {
+    return Call(op->op, get_analogous_expression(op->args), op->attrs, op->type_args, op->span);
+  }
+  Expr VisitExpr_(const LetNode* op) {
+    return Let(op->var, get_analogous_expression(op->value), get_analogous_expression(op->body),
+               op->span);
+  }
+  Expr VisitExpr_(const IfNode* op) {
+    return If(get_analogous_expression(op->cond), get_analogous_expression(op->true_branch),
+              get_analogous_expression(op->false_branch), op->span);
+  }
+  Expr VisitExpr_(const TupleGetItemNode* op) {
+    return TupleGetItem(get_analogous_expression(op->tuple), op->index, op->span);
+  }
+  Expr VisitExpr_(const RefCreateNode* op) {
+    return RefCreate(get_analogous_expression(op->value), op->span);
+  }
+  Expr VisitExpr_(const RefReadNode* op) {
+    return RefRead(get_analogous_expression(op->ref), op->span);
+  }
+  Expr VisitExpr_(const RefWriteNode* op) {
+    return RefWrite(get_analogous_expression(op->ref), get_analogous_expression(op->value),
+                    op->span);
+  }
+  Expr VisitExpr_(const ConstructorNode* op) {
+    return Constructor(op->name_hint, op->inputs, op->belong_to);
+  }
+  Expr VisitExpr_(const MatchNode* op) {
+    return Match(get_analogous_expression(op->data), op->clauses, op->complete, op->span);
+  }
+
+ private:
+  Expr get_analogous_expression(const Expr& expr) {
+    if (!expr->checked_type_.defined()) {
+      return VisitExpr(expr);
+    }
+
+    return Var("dummy_var", expr->checked_type(), expr->span);
+  }
+  Array<Expr> get_analogous_expression(const Array<Expr>& fields) {
+    Array<Expr> new_fields;
+    for (Expr expr : fields) {
+      new_fields.push_back(get_analogous_expression(expr));
+    }
+    return new_fields;
+  }
+};
+
 // A callable which hashes std::pair
 struct pair_hash {
   template <class T1, class T2>
@@ -194,23 +261,13 @@ class MixedPrecisionPass : public MixedModeMutator {
   }
 
   Type GetType(const Expr& expr) const {
-    const Type old_checked_type = expr->checked_type_;
-    if (old_checked_type.defined()) {
-      return old_checked_type;
-    }
-
-    auto mod = IRModule::FromExpr(MakeAnalogousSubgraph(expr));
-    // LOG(WARNING) << mod;
-    // LOG(WARNING) << IRModule::FromExpr(expr);
-    mod = transform::InferType()(mod);
-    Type t;
-    if (expr.as<FunctionNode>()) {
-      t = mod->Lookup("main")->checked_type();
-    } else {
-      t = mod->Lookup("main").as<FunctionNode>()->body->checked_type();
+    Type checked_type = expr->checked_type_;
+    if (checked_type.defined()) {
+      return checked_type;
     }
-    expr->checked_type_ = t;
-    return t;
+    checked_type = transform::InferTypeFast(expr);
+    expr->checked_type_ = checked_type;
+    return checked_type;
   }
 
   bool IsMixedPrecisionType(const Type& t, bool ignore_non_float = false) const {
diff --git a/src/relay/transforms/type_infer.cc b/src/relay/transforms/type_infer.cc
index 22bc8f34149c..03a5be7e6120 100644
--- a/src/relay/transforms/type_infer.cc
+++ b/src/relay/transforms/type_infer.cc
@@ -824,8 +824,90 @@ void AddGlobalTypes(IRModule mod) {
   }
 }
 
+class SameTypedSubgraphExtractor : public ExprMutator {
+  /*
+  Creates a small subgraph with the same type as the input expression.
+
+  ExprMutator is sufficient over MixedModemutator since we will not recurse much.
+  */
+
+  Expr VisitExpr_(const VarNode* op) { return Var(op->vid, op->type_annotation, op->span); }
+  Expr VisitExpr_(const ConstantNode* op) { return Constant(op->data, op->span); }
+  Expr VisitExpr_(const GlobalVarNode* op) { return GlobalVar(op->name_hint); }
+  Expr VisitExpr_(const OpNode* op) { return Op(GetRef<Op>(op)); }
+  Expr VisitExpr_(const TupleNode* op) {
+    return Tuple(get_analogous_expression(op->fields), op->span);
+  }
+  Expr VisitExpr_(const FunctionNode* op) {
+    // Here will be the only VisitExpr
+    return Function(op->params, get_analogous_expression(op->body), op->ret_type, op->type_params,
+                    op->attrs, op->span);
+  }
+  Expr VisitExpr_(const CallNode* op) {
+    return Call(op->op, get_analogous_expression(op->args), op->attrs, op->type_args, op->span);
+  }
+  Expr VisitExpr_(const LetNode* op) {
+    return Let(op->var, get_analogous_expression(op->value), get_analogous_expression(op->body),
+               op->span);
+  }
+  Expr VisitExpr_(const IfNode* op) {
+    return If(get_analogous_expression(op->cond), get_analogous_expression(op->true_branch),
+              get_analogous_expression(op->false_branch), op->span);
+  }
+  Expr VisitExpr_(const TupleGetItemNode* op) {
+    return TupleGetItem(get_analogous_expression(op->tuple), op->index, op->span);
+  }
+  Expr VisitExpr_(const RefCreateNode* op) {
+    return RefCreate(get_analogous_expression(op->value), op->span);
+  }
+  Expr VisitExpr_(const RefReadNode* op) {
+    return RefRead(get_analogous_expression(op->ref), op->span);
+  }
+  Expr VisitExpr_(const RefWriteNode* op) {
+    return RefWrite(get_analogous_expression(op->ref), get_analogous_expression(op->value),
+                    op->span);
+  }
+  Expr VisitExpr_(const ConstructorNode* op) {
+    return Constructor(op->name_hint, op->inputs, op->belong_to);
+  }
+  Expr VisitExpr_(const MatchNode* op) {
+    return Match(get_analogous_expression(op->data), op->clauses, op->complete, op->span);
+  }
+
+ private:
+  Expr get_analogous_expression(const Expr& expr) {
+    if (!expr->checked_type_.defined()) {
+      return VisitExpr(expr);
+    }
+
+    return Var("dummy_var", expr->checked_type(), expr->span);
+  }
+  Array<Expr> get_analogous_expression(const Array<Expr>& fields) {
+    Array<Expr> new_fields;
+    for (Expr expr : fields) {
+      new_fields.push_back(get_analogous_expression(expr));
+    }
+    return new_fields;
+  }
+};
+
 namespace transform {
 
+Type InferTypeFast(const Expr& expr) {
+  SameTypedSubgraphExtractor subgraph_extractor;
+  auto mod = IRModule::FromExpr(subgraph_extractor(expr));
+  mod = transform::InferType()(mod);
+  if (expr.as<FunctionNode>()) {
+    return mod->Lookup("main")->checked_type();
+  } else {
+    return mod->Lookup("main").as<FunctionNode>()->body->checked_type();
+  }
+}
+
+TVM_REGISTER_GLOBAL("relay._transform.InferTypeFast").set_body_typed([](const Expr& expr) {
+  return InferTypeFast(expr);
+});
+
 Pass InferType() {
   auto pass_info = PassInfo(0, "InferType", {});
   return tvm::transform::CreateModulePass(

From 7dee27bcb8b1ec36db7fbbaf36281b64309cd1b9 Mon Sep 17 00:00:00 2001
From: Andrew Luo <andrew.zhao.luo@gmail.com>
Date: Mon, 13 Dec 2021 16:59:50 -0800
Subject: [PATCH 04/19] clean up src logs

---
 src/relay/op/nn/nn.h | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/relay/op/nn/nn.h b/src/relay/op/nn/nn.h
index 900c61eb8200..6bc21473af18 100644
--- a/src/relay/op/nn/nn.h
+++ b/src/relay/op/nn/nn.h
@@ -162,7 +162,6 @@ bool BatchMatmulRel(const Array<Type>& types, int num_inputs, const Attrs& attrs
   if (out_dtype.bits() == 0) {
     out_dtype = x->dtype;
   }
-
   // assign output type
   const auto& out_b =
       xb->IsInstance<tir::AnyNode>() || yb->IsInstance<tir::AnyNode>() ? tir::Any() : max(xb, yb);

From 08b391ae9602725781c60f536302b1844f3d4dc1 Mon Sep 17 00:00:00 2001
From: Andrew Luo <andrew.zhao.luo@gmail.com>
Date: Mon, 13 Dec 2021 17:01:01 -0800
Subject: [PATCH 05/19] clean up PR more

---
 src/relay/transforms/to_mixed_precision.cc | 70 ----------------------
 1 file changed, 70 deletions(-)

diff --git a/src/relay/transforms/to_mixed_precision.cc b/src/relay/transforms/to_mixed_precision.cc
index ef8a78c330aa..9d0804d06623 100644
--- a/src/relay/transforms/to_mixed_precision.cc
+++ b/src/relay/transforms/to_mixed_precision.cc
@@ -36,73 +36,6 @@
 namespace tvm {
 namespace relay {
 
-class SameTypedSubgraphExtractor : public ExprMutator {
-  /*
-  Creates a small subgraph with the same type as the input expression.
-
-  ExprMutator is sufficient over MixedModemutator since we will not recurse much.
-  */
-
-  Expr VisitExpr_(const VarNode* op) { return Var(op->vid, op->type_annotation, op->span); }
-  Expr VisitExpr_(const ConstantNode* op) { return Constant(op->data, op->span); }
-  Expr VisitExpr_(const GlobalVarNode* op) { return GlobalVar(op->name_hint); }
-  Expr VisitExpr_(const OpNode* op) { return Op(GetRef<Op>(op)); }
-  Expr VisitExpr_(const TupleNode* op) {
-    return Tuple(get_analogous_expression(op->fields), op->span);
-  }
-  Expr VisitExpr_(const FunctionNode* op) {
-    // Here will be the only VisitExpr
-    return Function(op->params, get_analogous_expression(op->body), op->ret_type, op->type_params,
-                    op->attrs, op->span);
-  }
-  Expr VisitExpr_(const CallNode* op) {
-    return Call(op->op, get_analogous_expression(op->args), op->attrs, op->type_args, op->span);
-  }
-  Expr VisitExpr_(const LetNode* op) {
-    return Let(op->var, get_analogous_expression(op->value), get_analogous_expression(op->body),
-               op->span);
-  }
-  Expr VisitExpr_(const IfNode* op) {
-    return If(get_analogous_expression(op->cond), get_analogous_expression(op->true_branch),
-              get_analogous_expression(op->false_branch), op->span);
-  }
-  Expr VisitExpr_(const TupleGetItemNode* op) {
-    return TupleGetItem(get_analogous_expression(op->tuple), op->index, op->span);
-  }
-  Expr VisitExpr_(const RefCreateNode* op) {
-    return RefCreate(get_analogous_expression(op->value), op->span);
-  }
-  Expr VisitExpr_(const RefReadNode* op) {
-    return RefRead(get_analogous_expression(op->ref), op->span);
-  }
-  Expr VisitExpr_(const RefWriteNode* op) {
-    return RefWrite(get_analogous_expression(op->ref), get_analogous_expression(op->value),
-                    op->span);
-  }
-  Expr VisitExpr_(const ConstructorNode* op) {
-    return Constructor(op->name_hint, op->inputs, op->belong_to);
-  }
-  Expr VisitExpr_(const MatchNode* op) {
-    return Match(get_analogous_expression(op->data), op->clauses, op->complete, op->span);
-  }
-
- private:
-  Expr get_analogous_expression(const Expr& expr) {
-    if (!expr->checked_type_.defined()) {
-      return VisitExpr(expr);
-    }
-
-    return Var("dummy_var", expr->checked_type(), expr->span);
-  }
-  Array<Expr> get_analogous_expression(const Array<Expr>& fields) {
-    Array<Expr> new_fields;
-    for (Expr expr : fields) {
-      new_fields.push_back(get_analogous_expression(expr));
-    }
-    return new_fields;
-  }
-};
-
 // A callable which hashes std::pair
 struct pair_hash {
   template <class T1, class T2>
@@ -168,9 +101,6 @@ class MixedPrecisionPass : public MixedModeMutator {
   /*! \brief The target datatype we want to convert to e.g. FP16 */
   const DataType mixed_precision_type_;
 
-  /* TODO*/
-  std::unordered_map<Expr, Expr, ObjectPtrHash, ObjectPtrEqual> analgous_graphs;
-
   /*! \brief Map of Ops with no associated FTVMMixedPrecisionConversionType to the times they were
    * encountered. Used for emitting warnings on missing ops in the pass.
    */

From 2021f23110e1d2baf4a90260e672afc678314594 Mon Sep 17 00:00:00 2001
From: Andrew Luo <andrew.zhao.luo@gmail.com>
Date: Mon, 13 Dec 2021 17:01:55 -0800
Subject: [PATCH 06/19] more clean up

---
 src/relay/transforms/to_mixed_precision.cc | 15 ---------------
 1 file changed, 15 deletions(-)

diff --git a/src/relay/transforms/to_mixed_precision.cc b/src/relay/transforms/to_mixed_precision.cc
index 9d0804d06623..02726060c90f 100644
--- a/src/relay/transforms/to_mixed_precision.cc
+++ b/src/relay/transforms/to_mixed_precision.cc
@@ -175,21 +175,6 @@ class MixedPrecisionPass : public MixedModeMutator {
     return Attrs(new_attrs);
   }
 
-  Expr MakeAnalogousSubgraph(const Expr& expr) const {
-    if (auto node = expr.as<CallNode>()) {
-      Array<Expr> args;
-      for (Expr expr : node->args) {
-        args.push_back(Var("dummy_temp", GetType(expr)));
-      }
-      return Call(node->op, args, node->attrs, node->type_args, node->span);
-    } else if (auto node = expr.as<TupleGetItemNode>()) {
-      return TupleGetItem(MakeAnalogousSubgraph(node->tuple), node->index, node->span);
-    } else {
-      LOG(FATAL) << "Unknown node " << expr;
-      return Expr(nullptr);
-    }
-  }
-
   Type GetType(const Expr& expr) const {
     Type checked_type = expr->checked_type_;
     if (checked_type.defined()) {

From 5136b85a30a509339eb88313a6e9493211a18e99 Mon Sep 17 00:00:00 2001
From: Andrew Luo <andrew.zhao.luo@gmail.com>
Date: Tue, 14 Dec 2021 15:08:03 -0800
Subject: [PATCH 07/19] more documenetation

---
 include/tvm/relay/transform.h      | 15 +++++++++++----
 src/relay/transforms/type_infer.cc |  7 ++++++-
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/include/tvm/relay/transform.h b/include/tvm/relay/transform.h
index a0624a890858..800b0839ef0b 100644
--- a/include/tvm/relay/transform.h
+++ b/include/tvm/relay/transform.h
@@ -250,7 +250,7 @@ TVM_DLL Pass DynamicToStatic();
 /*!
  * \brief Infer the type of an expression.
  *
- * The result of type checking is a new expression with unambigous
+ * The result of type checking is a new expression with unambiguous
  * type information filled in, as well as it's checked type field
  * populated with the result type.
  *
@@ -258,9 +258,16 @@ TVM_DLL Pass DynamicToStatic();
  */
 TVM_DLL Pass InferType();
 
-/*
-TODO
-*/
+/*!
+ * \brief Infer the type of an expression.
+ *
+ * The result of type checking is a new expression with unambiguous
+ * type information filled in for that expression only. The fast
+ * version depends on existing type information populated throughout
+ * the expression and assumes this information is correct.
+ *
+ * \return The pass.
+ */
 TVM_DLL Type InferTypeFast(const Expr& expr);
 
 /*!
diff --git a/src/relay/transforms/type_infer.cc b/src/relay/transforms/type_infer.cc
index 03a5be7e6120..3b823f73c34d 100644
--- a/src/relay/transforms/type_infer.cc
+++ b/src/relay/transforms/type_infer.cc
@@ -826,7 +826,11 @@ void AddGlobalTypes(IRModule mod) {
 
 class SameTypedSubgraphExtractor : public ExprMutator {
   /*
-  Creates a small subgraph with the same type as the input expression.
+  Creates a small subgraph with the same type as the input expression. We attempt to do
+  by depending on existing type information being populated in expressions the target
+  node depends on. If a node with populated type information is found we simply
+  replace it with a variable of that type. In this way, we can avoid copying and
+  recursing through most of the expression graph.
 
   ExprMutator is sufficient over MixedModemutator since we will not recurse much.
   */
@@ -876,6 +880,7 @@ class SameTypedSubgraphExtractor : public ExprMutator {
 
  private:
   Expr get_analogous_expression(const Expr& expr) {
+    // Replace the expression with a potentially simpler expression of the same type
     if (!expr->checked_type_.defined()) {
       return VisitExpr(expr);
     }

From dbf3cf6fce4f4691611154c49f50af73c57e1273 Mon Sep 17 00:00:00 2001
From: Andrew Luo <andrew.zhao.luo@gmail.com>
Date: Tue, 14 Dec 2021 15:19:26 -0800
Subject: [PATCH 08/19] clean up

---
 include/tvm/relay/transform.h              |  6 ++++--
 src/relay/transforms/to_mixed_precision.cc |  6 +++---
 src/relay/transforms/type_infer.cc         | 10 +++++-----
 3 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/include/tvm/relay/transform.h b/include/tvm/relay/transform.h
index 800b0839ef0b..a877860360ae 100644
--- a/include/tvm/relay/transform.h
+++ b/include/tvm/relay/transform.h
@@ -259,12 +259,14 @@ TVM_DLL Pass DynamicToStatic();
 TVM_DLL Pass InferType();
 
 /*!
- * \brief Infer the type of an expression.
+ * \brief Infer the type of an expression, reusing existing type information.
  *
  * The result of type checking is a new expression with unambiguous
  * type information filled in for that expression only. The fast
  * version depends on existing type information populated throughout
- * the expression and assumes this information is correct.
+ * the expression and assumes this information is correct. The fast 
+ * version also avoids examining large amounts of the graph assuming
+ * type information is filled in.
  *
  * \return The pass.
  */
diff --git a/src/relay/transforms/to_mixed_precision.cc b/src/relay/transforms/to_mixed_precision.cc
index 02726060c90f..faa706b99ea0 100644
--- a/src/relay/transforms/to_mixed_precision.cc
+++ b/src/relay/transforms/to_mixed_precision.cc
@@ -180,9 +180,9 @@ class MixedPrecisionPass : public MixedModeMutator {
     if (checked_type.defined()) {
       return checked_type;
     }
-    checked_type = transform::InferTypeFast(expr);
-    expr->checked_type_ = checked_type;
-    return checked_type;
+
+    // This also populates the checked_type_ field for expr
+    return transform::InferTypeFast(expr);
   }
 
   bool IsMixedPrecisionType(const Type& t, bool ignore_non_float = false) const {
diff --git a/src/relay/transforms/type_infer.cc b/src/relay/transforms/type_infer.cc
index 3b823f73c34d..1035235cd128 100644
--- a/src/relay/transforms/type_infer.cc
+++ b/src/relay/transforms/type_infer.cc
@@ -901,12 +901,12 @@ namespace transform {
 Type InferTypeFast(const Expr& expr) {
   SameTypedSubgraphExtractor subgraph_extractor;
   auto mod = IRModule::FromExpr(subgraph_extractor(expr));
+
   mod = transform::InferType()(mod);
-  if (expr.as<FunctionNode>()) {
-    return mod->Lookup("main")->checked_type();
-  } else {
-    return mod->Lookup("main").as<FunctionNode>()->body->checked_type();
-  }
+  Type result_type = mod->Lookup("main").as<FunctionNode>()->body->checked_type();
+
+  expr->checked_type_ = result_type;
+  return result_type;
 }
 
 TVM_REGISTER_GLOBAL("relay._transform.InferTypeFast").set_body_typed([](const Expr& expr) {

From e9a5f55c06879977a384b464dbec7e3da908b2c6 Mon Sep 17 00:00:00 2001
From: Andrew Luo <andrew.zhao.luo@gmail.com>
Date: Tue, 14 Dec 2021 15:37:55 -0800
Subject: [PATCH 09/19] formatting

---
 include/tvm/relay/transform.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/include/tvm/relay/transform.h b/include/tvm/relay/transform.h
index a877860360ae..b0cb8d6decc0 100644
--- a/include/tvm/relay/transform.h
+++ b/include/tvm/relay/transform.h
@@ -264,7 +264,7 @@ TVM_DLL Pass InferType();
  * The result of type checking is a new expression with unambiguous
  * type information filled in for that expression only. The fast
  * version depends on existing type information populated throughout
- * the expression and assumes this information is correct. The fast 
+ * the expression and assumes this information is correct. The fast
  * version also avoids examining large amounts of the graph assuming
  * type information is filled in.
  *

From f8c5012a48762a452851ba4dc6542d5afd7dcf35 Mon Sep 17 00:00:00 2001
From: Andrew Luo <andrew.zhao.luo@gmail.com>
Date: Wed, 15 Dec 2021 10:39:43 -0800
Subject: [PATCH 10/19] rename fast --> local

---
 include/tvm/relay/transform.h              | 9 +++++----
 src/relay/transforms/to_mixed_precision.cc | 2 +-
 src/relay/transforms/type_infer.cc         | 6 +++---
 3 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/include/tvm/relay/transform.h b/include/tvm/relay/transform.h
index b0cb8d6decc0..91d5b77eebf9 100644
--- a/include/tvm/relay/transform.h
+++ b/include/tvm/relay/transform.h
@@ -262,15 +262,16 @@ TVM_DLL Pass InferType();
  * \brief Infer the type of an expression, reusing existing type information.
  *
  * The result of type checking is a new expression with unambiguous
- * type information filled in for that expression only. The fast
+ * type information filled in for that expression only. The local
  * version depends on existing type information populated throughout
- * the expression and assumes this information is correct. The fast
+ * the expression and assumes this information is correct. The local
  * version also avoids examining large amounts of the graph assuming
- * type information is filled in.
+ * type information is filled in properly which makes it much faster if we
+ * iteratively call type inference.
  *
  * \return The pass.
  */
-TVM_DLL Type InferTypeFast(const Expr& expr);
+TVM_DLL Type InferTypeLocal(const Expr& expr);
 
 /*!
  * \brief Search and eliminate common subexpression. For example, if there are
diff --git a/src/relay/transforms/to_mixed_precision.cc b/src/relay/transforms/to_mixed_precision.cc
index faa706b99ea0..6155ec7787c9 100644
--- a/src/relay/transforms/to_mixed_precision.cc
+++ b/src/relay/transforms/to_mixed_precision.cc
@@ -182,7 +182,7 @@ class MixedPrecisionPass : public MixedModeMutator {
     }
 
     // This also populates the checked_type_ field for expr
-    return transform::InferTypeFast(expr);
+    return transform::InferTypeLocal(expr);
   }
 
   bool IsMixedPrecisionType(const Type& t, bool ignore_non_float = false) const {
diff --git a/src/relay/transforms/type_infer.cc b/src/relay/transforms/type_infer.cc
index 1035235cd128..a3f68c75cddf 100644
--- a/src/relay/transforms/type_infer.cc
+++ b/src/relay/transforms/type_infer.cc
@@ -898,7 +898,7 @@ class SameTypedSubgraphExtractor : public ExprMutator {
 
 namespace transform {
 
-Type InferTypeFast(const Expr& expr) {
+Type InferTypeLocal(const Expr& expr) {
   SameTypedSubgraphExtractor subgraph_extractor;
   auto mod = IRModule::FromExpr(subgraph_extractor(expr));
 
@@ -909,8 +909,8 @@ Type InferTypeFast(const Expr& expr) {
   return result_type;
 }
 
-TVM_REGISTER_GLOBAL("relay._transform.InferTypeFast").set_body_typed([](const Expr& expr) {
-  return InferTypeFast(expr);
+TVM_REGISTER_GLOBAL("relay._transform.InferTypeLocal").set_body_typed([](const Expr& expr) {
+  return InferTypeLocal(expr);
 });
 
 Pass InferType() {

From 5960c5c8edd99d95bfcbc1eedf46a35cecb146a5 Mon Sep 17 00:00:00 2001
From: Andrew Luo <andrew.zhao.luo@gmail.com>
Date: Wed, 15 Dec 2021 10:52:06 -0800
Subject: [PATCH 11/19] more ocmments

---
 include/tvm/relay/transform.h      |  4 ++--
 src/relay/transforms/type_infer.cc | 14 +++++++++++++-
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/include/tvm/relay/transform.h b/include/tvm/relay/transform.h
index 91d5b77eebf9..78dda45b954a 100644
--- a/include/tvm/relay/transform.h
+++ b/include/tvm/relay/transform.h
@@ -262,8 +262,8 @@ TVM_DLL Pass InferType();
  * \brief Infer the type of an expression, reusing existing type information.
  *
  * The result of type checking is a new expression with unambiguous
- * type information filled in for that expression only. The local
- * version depends on existing type information populated throughout
+ * type information filled in for the given node only. The local
+ * version can use existing type information populated throughout
  * the expression and assumes this information is correct. The local
  * version also avoids examining large amounts of the graph assuming
  * type information is filled in properly which makes it much faster if we
diff --git a/src/relay/transforms/type_infer.cc b/src/relay/transforms/type_infer.cc
index a3f68c75cddf..5c3cc4c16f6c 100644
--- a/src/relay/transforms/type_infer.cc
+++ b/src/relay/transforms/type_infer.cc
@@ -830,7 +830,8 @@ class SameTypedSubgraphExtractor : public ExprMutator {
   by depending on existing type information being populated in expressions the target
   node depends on. If a node with populated type information is found we simply
   replace it with a variable of that type. In this way, we can avoid copying and
-  recursing through most of the expression graph.
+  recursing through most of the expression graph. Note, this assumes that current
+  populated type information is correct!
 
   ExprMutator is sufficient over MixedModemutator since we will not recurse much.
   */
@@ -899,6 +900,17 @@ class SameTypedSubgraphExtractor : public ExprMutator {
 namespace transform {
 
 Type InferTypeLocal(const Expr& expr) {
+  /*
+  This type inference differs from InferType in that it uses existing type information
+  to avoid recursing over much of the graph, and it only examines the type of the input
+  node. This makes it faster if you need to run type inference iteratively throughout
+  a pass for example.
+
+  However, it assumes any existing populated type inference is correct! If some populated
+  type inference is incorrect, an incorrect type may be returned or a type error will be
+  raised. If you know not all populated type fields are correct with the current graph,
+  you should use InferType() instead.
+  */
   SameTypedSubgraphExtractor subgraph_extractor;
   auto mod = IRModule::FromExpr(subgraph_extractor(expr));
 

From f294f6385eda91e36b9e3d20c7898d917eb493aa Mon Sep 17 00:00:00 2001
From: Andrew Luo <andrew.zhao.luo@gmail.com>
Date: Thu, 16 Dec 2021 16:10:21 -0800
Subject: [PATCH 12/19] jostle ci


From 4f0b03bceb784c71c862ab83f78f029cfacf9a91 Mon Sep 17 00:00:00 2001
From: Andrew Luo <andrew.zhao.luo@gmail.com>
Date: Fri, 17 Dec 2021 11:04:00 -0800
Subject: [PATCH 13/19] type inference

---
 src/relay/transforms/type_infer.cc | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/src/relay/transforms/type_infer.cc b/src/relay/transforms/type_infer.cc
index 5c3cc4c16f6c..9756592cd9e5 100644
--- a/src/relay/transforms/type_infer.cc
+++ b/src/relay/transforms/type_infer.cc
@@ -844,9 +844,13 @@ class SameTypedSubgraphExtractor : public ExprMutator {
     return Tuple(get_analogous_expression(op->fields), op->span);
   }
   Expr VisitExpr_(const FunctionNode* op) {
-    // Here will be the only VisitExpr
-    return Function(op->params, get_analogous_expression(op->body), op->ret_type, op->type_params,
-                    op->attrs, op->span);
+    // We use these to regenerate the list of free variables in the function and place them in
+    // the list of input parameters for the model.
+    Expr new_body = get_analogous_expression(op->body);
+    IRModule new_body_mod = IRModule::FromExpr(new_body);
+    return Function(relay::FreeVars(new_body), new_body, op->ret_type,
+                    relay::FreeTypeVars(new_body, IRModule::FromExpr(new_body)), op->attrs,
+                    op->span);
   }
   Expr VisitExpr_(const CallNode* op) {
     return Call(op->op, get_analogous_expression(op->args), op->attrs, op->type_args, op->span);

From 830105719830017c8e8aafbd5d5804754f021288 Mon Sep 17 00:00:00 2001
From: Andrew Luo <andrew.zhao.luo@gmail.com>
Date: Fri, 17 Dec 2021 11:10:24 -0800
Subject: [PATCH 14/19] change comment for SameTypedSubgraphExtractor

---
 src/relay/transforms/type_infer.cc | 22 +++++++++++-----------
 1 file changed, 11 insertions(+), 11 deletions(-)

diff --git a/src/relay/transforms/type_infer.cc b/src/relay/transforms/type_infer.cc
index 9756592cd9e5..5289ffbdb14c 100644
--- a/src/relay/transforms/type_infer.cc
+++ b/src/relay/transforms/type_infer.cc
@@ -824,18 +824,18 @@ void AddGlobalTypes(IRModule mod) {
   }
 }
 
+/*!
+ * \brief Returns a possibly much smaller subgraph whose inner nodes have the same type.
+ *
+ * Returns the largest sub-graph who's inner nodes need types and leaves are vars standing in
+ * for already typed sub-expressions. This creates a graph whose inner nodes have the same
+ * type as the original graph and when running type inference, we can avoid copying and
+ * recursing through most of the expression graph when running type inference. Note, this assumes
+ * that current populated type information is correct!
+ *
+ * ExprMutator is sufficient over MixedModemutator since we will not recurse much.
+ */
 class SameTypedSubgraphExtractor : public ExprMutator {
-  /*
-  Creates a small subgraph with the same type as the input expression. We attempt to do
-  by depending on existing type information being populated in expressions the target
-  node depends on. If a node with populated type information is found we simply
-  replace it with a variable of that type. In this way, we can avoid copying and
-  recursing through most of the expression graph. Note, this assumes that current
-  populated type information is correct!
-
-  ExprMutator is sufficient over MixedModemutator since we will not recurse much.
-  */
-
   Expr VisitExpr_(const VarNode* op) { return Var(op->vid, op->type_annotation, op->span); }
   Expr VisitExpr_(const ConstantNode* op) { return Constant(op->data, op->span); }
   Expr VisitExpr_(const GlobalVarNode* op) { return GlobalVar(op->name_hint); }

From 1cb38f11c9d84fab60466e3ffb8e427da96b4ffc Mon Sep 17 00:00:00 2001
From: Andrew Luo <andrew.zhao.luo@gmail.com>
Date: Fri, 17 Dec 2021 11:11:08 -0800
Subject: [PATCH 15/19] get_analogous_expression -> GetAnalogousExpression

---
 src/relay/transforms/type_infer.cc | 29 ++++++++++++++---------------
 1 file changed, 14 insertions(+), 15 deletions(-)

diff --git a/src/relay/transforms/type_infer.cc b/src/relay/transforms/type_infer.cc
index 5289ffbdb14c..9dea45414fa4 100644
--- a/src/relay/transforms/type_infer.cc
+++ b/src/relay/transforms/type_infer.cc
@@ -841,50 +841,49 @@ class SameTypedSubgraphExtractor : public ExprMutator {
   Expr VisitExpr_(const GlobalVarNode* op) { return GlobalVar(op->name_hint); }
   Expr VisitExpr_(const OpNode* op) { return Op(GetRef<Op>(op)); }
   Expr VisitExpr_(const TupleNode* op) {
-    return Tuple(get_analogous_expression(op->fields), op->span);
+    return Tuple(GetAnalogousExpression(op->fields), op->span);
   }
   Expr VisitExpr_(const FunctionNode* op) {
     // We use these to regenerate the list of free variables in the function and place them in
     // the list of input parameters for the model.
-    Expr new_body = get_analogous_expression(op->body);
+    Expr new_body = GetAnalogousExpression(op->body);
     IRModule new_body_mod = IRModule::FromExpr(new_body);
     return Function(relay::FreeVars(new_body), new_body, op->ret_type,
                     relay::FreeTypeVars(new_body, IRModule::FromExpr(new_body)), op->attrs,
                     op->span);
   }
   Expr VisitExpr_(const CallNode* op) {
-    return Call(op->op, get_analogous_expression(op->args), op->attrs, op->type_args, op->span);
+    return Call(op->op, GetAnalogousExpression(op->args), op->attrs, op->type_args, op->span);
   }
   Expr VisitExpr_(const LetNode* op) {
-    return Let(op->var, get_analogous_expression(op->value), get_analogous_expression(op->body),
+    return Let(op->var, GetAnalogousExpression(op->value), GetAnalogousExpression(op->body),
                op->span);
   }
   Expr VisitExpr_(const IfNode* op) {
-    return If(get_analogous_expression(op->cond), get_analogous_expression(op->true_branch),
-              get_analogous_expression(op->false_branch), op->span);
+    return If(GetAnalogousExpression(op->cond), GetAnalogousExpression(op->true_branch),
+              GetAnalogousExpression(op->false_branch), op->span);
   }
   Expr VisitExpr_(const TupleGetItemNode* op) {
-    return TupleGetItem(get_analogous_expression(op->tuple), op->index, op->span);
+    return TupleGetItem(GetAnalogousExpression(op->tuple), op->index, op->span);
   }
   Expr VisitExpr_(const RefCreateNode* op) {
-    return RefCreate(get_analogous_expression(op->value), op->span);
+    return RefCreate(GetAnalogousExpression(op->value), op->span);
   }
   Expr VisitExpr_(const RefReadNode* op) {
-    return RefRead(get_analogous_expression(op->ref), op->span);
+    return RefRead(GetAnalogousExpression(op->ref), op->span);
   }
   Expr VisitExpr_(const RefWriteNode* op) {
-    return RefWrite(get_analogous_expression(op->ref), get_analogous_expression(op->value),
-                    op->span);
+    return RefWrite(GetAnalogousExpression(op->ref), GetAnalogousExpression(op->value), op->span);
   }
   Expr VisitExpr_(const ConstructorNode* op) {
     return Constructor(op->name_hint, op->inputs, op->belong_to);
   }
   Expr VisitExpr_(const MatchNode* op) {
-    return Match(get_analogous_expression(op->data), op->clauses, op->complete, op->span);
+    return Match(GetAnalogousExpression(op->data), op->clauses, op->complete, op->span);
   }
 
  private:
-  Expr get_analogous_expression(const Expr& expr) {
+  Expr GetAnalogousExpression(const Expr& expr) {
     // Replace the expression with a potentially simpler expression of the same type
     if (!expr->checked_type_.defined()) {
       return VisitExpr(expr);
@@ -892,10 +891,10 @@ class SameTypedSubgraphExtractor : public ExprMutator {
 
     return Var("dummy_var", expr->checked_type(), expr->span);
   }
-  Array<Expr> get_analogous_expression(const Array<Expr>& fields) {
+  Array<Expr> GetAnalogousExpression(const Array<Expr>& fields) {
     Array<Expr> new_fields;
     for (Expr expr : fields) {
-      new_fields.push_back(get_analogous_expression(expr));
+      new_fields.push_back(GetAnalogousExpression(expr));
     }
     return new_fields;
   }

From 5aae167b3f825b3fccafa910e245c22df6b33cc7 Mon Sep 17 00:00:00 2001
From: Andrew Luo <andrew.zhao.luo@gmail.com>
Date: Fri, 17 Dec 2021 11:12:37 -0800
Subject: [PATCH 16/19] comment in GetAnaalogousExpression

---
 src/relay/transforms/type_infer.cc | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/relay/transforms/type_infer.cc b/src/relay/transforms/type_infer.cc
index 9dea45414fa4..e3a14fcef936 100644
--- a/src/relay/transforms/type_infer.cc
+++ b/src/relay/transforms/type_infer.cc
@@ -885,11 +885,13 @@ class SameTypedSubgraphExtractor : public ExprMutator {
  private:
   Expr GetAnalogousExpression(const Expr& expr) {
     // Replace the expression with a potentially simpler expression of the same type
-    if (!expr->checked_type_.defined()) {
-      return VisitExpr(expr);
+    if (expr->checked_type_.defined()) {
+      // Since the expression already has a checked_type which we assume is correct we don't need
+      // full type inference to enter it. So stub it out with a dummy var of the same type.
+      return Var("dummy_var", expr->checked_type(), expr->span);
     }
 
-    return Var("dummy_var", expr->checked_type(), expr->span);
+    return VisitExpr(expr);
   }
   Array<Expr> GetAnalogousExpression(const Array<Expr>& fields) {
     Array<Expr> new_fields;

From d6f73f27ec908113e7670704c728ee417ea2ef33 Mon Sep 17 00:00:00 2001
From: Andrew Luo <andrew.zhao.luo@gmail.com>
Date: Fri, 17 Dec 2021 11:31:28 -0800
Subject: [PATCH 17/19] add comment

---
 src/relay/transforms/to_mixed_precision.cc | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/relay/transforms/to_mixed_precision.cc b/src/relay/transforms/to_mixed_precision.cc
index 6155ec7787c9..d8c7aa2ffcfa 100644
--- a/src/relay/transforms/to_mixed_precision.cc
+++ b/src/relay/transforms/to_mixed_precision.cc
@@ -176,6 +176,10 @@ class MixedPrecisionPass : public MixedModeMutator {
   }
 
   Type GetType(const Expr& expr) const {
+    // The expression has not been changed AND it's existing type
+    // is known to still be valid. (See special handling for tuples etc
+    // below for where we null out checked_type_ when we can not
+    // sure it is still valid.
     Type checked_type = expr->checked_type_;
     if (checked_type.defined()) {
       return checked_type;

From 09fbbe024f64e76a8ae1627bdef1094fb0a1c066 Mon Sep 17 00:00:00 2001
From: Andrew Luo <andrew.zhao.luo@gmail.com>
Date: Fri, 17 Dec 2021 14:59:53 -0800
Subject: [PATCH 18/19] replace infer tests

---
 include/tvm/relay/transform.h           |  2 +-
 python/tvm/relay/transform/transform.py | 19 +++++++++++++++++
 src/relay/transforms/type_infer.cc      | 22 +++++++++++---------
 tests/python/relay/test_type_infer.py   | 27 ++++++++++++-------------
 4 files changed, 45 insertions(+), 25 deletions(-)

diff --git a/include/tvm/relay/transform.h b/include/tvm/relay/transform.h
index 78dda45b954a..ee5a19794577 100644
--- a/include/tvm/relay/transform.h
+++ b/include/tvm/relay/transform.h
@@ -269,7 +269,7 @@ TVM_DLL Pass InferType();
  * type information is filled in properly which makes it much faster if we
  * iteratively call type inference.
  *
- * \return The pass.
+ * \return The type of the expression.
  */
 TVM_DLL Type InferTypeLocal(const Expr& expr);
 
diff --git a/python/tvm/relay/transform/transform.py b/python/tvm/relay/transform/transform.py
index 4369009559ba..8613a4231e21 100644
--- a/python/tvm/relay/transform/transform.py
+++ b/python/tvm/relay/transform/transform.py
@@ -99,6 +99,25 @@ def InferType():
     return _ffi_api.InferType()
 
 
+def InferTypeLocal(expr):
+    """Infer the type of a single expr, reusing type information to do so.
+
+    This populates the checked_type field in expr. We assume existing type information
+    in the graph is correct!
+
+    Parameters
+    ----------
+    expr: relay.Expr
+        The expression we want to know the type of
+
+    Returns
+    -------
+    type: relay.Type
+        The type of the expression
+    """
+    return _ffi_api.InferTypeLocal(expr)
+
+
 def FoldScaleAxis():
     """Fold the scaling of axis into weights of conv2d/dense. This pass will
     invoke both forward and backward scale folding.
diff --git a/src/relay/transforms/type_infer.cc b/src/relay/transforms/type_infer.cc
index e3a14fcef936..456e210f7343 100644
--- a/src/relay/transforms/type_infer.cc
+++ b/src/relay/transforms/type_infer.cc
@@ -844,13 +844,9 @@ class SameTypedSubgraphExtractor : public ExprMutator {
     return Tuple(GetAnalogousExpression(op->fields), op->span);
   }
   Expr VisitExpr_(const FunctionNode* op) {
-    // We use these to regenerate the list of free variables in the function and place them in
-    // the list of input parameters for the model.
-    Expr new_body = GetAnalogousExpression(op->body);
-    IRModule new_body_mod = IRModule::FromExpr(new_body);
-    return Function(relay::FreeVars(new_body), new_body, op->ret_type,
-                    relay::FreeTypeVars(new_body, IRModule::FromExpr(new_body)), op->attrs,
-                    op->span);
+    // Unfortunately our strategy of inserting variables as dummies would change the signature of
+    // existing function nodes so we have to copy all used functions always :/
+    return Function(op->params, op->body, op->ret_type, op->type_params, op->attrs, op->span);
   }
   Expr VisitExpr_(const CallNode* op) {
     return Call(op->op, GetAnalogousExpression(op->args), op->attrs, op->type_args, op->span);
@@ -917,10 +913,16 @@ Type InferTypeLocal(const Expr& expr) {
   you should use InferType() instead.
   */
   SameTypedSubgraphExtractor subgraph_extractor;
-  auto mod = IRModule::FromExpr(subgraph_extractor(expr));
-
+  Expr sub_graph = subgraph_extractor(expr);
+  auto mod = IRModule::FromExpr(sub_graph);
   mod = transform::InferType()(mod);
-  Type result_type = mod->Lookup("main").as<FunctionNode>()->body->checked_type();
+
+  Type result_type;
+  if (expr.as<FunctionNode>()) {
+    result_type = mod->Lookup("main")->checked_type();
+  } else {
+    result_type = mod->Lookup("main").as<FunctionNode>()->body->checked_type();
+  }
 
   expr->checked_type_ = result_type;
   return result_type;
diff --git a/tests/python/relay/test_type_infer.py b/tests/python/relay/test_type_infer.py
index a0d37844b837..af64ce714df8 100644
--- a/tests/python/relay/test_type_infer.py
+++ b/tests/python/relay/test_type_infer.py
@@ -19,9 +19,8 @@
 """
 import pytest
 import tvm
-
-from tvm import IRModule, te, relay, parser
-from tvm.relay import op, transform, analysis
+from tvm import IRModule, parser, relay, te
+from tvm.relay import analysis, op, transform
 from tvm.relay.op import op as _op
 
 
@@ -33,12 +32,9 @@ def infer_mod(mod, annotate_spans=True):
     return mod
 
 
-def infer_expr(expr, annotate_spans=True):
-    mod = IRModule.from_expr(expr)
-    mod = infer_mod(mod, annotate_spans)
-    mod = transform.InferType()(mod)
-    entry = mod["main"]
-    return entry if isinstance(expr, relay.Function) else entry.body
+def infer_expr(expr):
+    transform.InferTypeLocal(expr)
+    return expr
 
 
 def assert_has_type(expr, typ, mod=None):
@@ -68,7 +64,7 @@ def test_monomorphic_let():
     # TODO(@jroesch): this seems whack.
     sb = relay.ScopeBuilder()
     x = relay.var("x", dtype="float64", shape=())
-    x = sb.let("x", relay.const(1.0, "float64"))
+    x = sb.let(x, relay.const(1.0, "float64"))
     sb.ret(x)
     xchecked = infer_expr(sb.get())
     assert xchecked.checked_type == relay.scalar_type("float64")
@@ -165,11 +161,11 @@ def @f(%n: int32, %data: float32) -> float32 {
 def test_incomplete_call():
     tt = relay.scalar_type("int32")
     x = relay.var("x", tt)
+    f_type = relay.FuncType([tt], tt)
     f = relay.var("f")
     func = relay.Function([x, f], relay.Call(f, [x]), tt)
 
     ft = infer_expr(func)
-    f_type = relay.FuncType([tt], tt)
     assert ft.checked_type == relay.FuncType([tt, f_type], tt)
 
 
@@ -245,7 +241,7 @@ def test_ref():
 def test_free_expr():
     x = relay.var("x", "float32")
     y = relay.add(x, x)
-    yy = infer_expr(y, annotate_spans=False)
+    yy = infer_expr(y)
     assert tvm.ir.structural_equal(yy.args[0], x, map_free_vars=True)
     assert yy.checked_type == relay.scalar_type("float32")
     assert x.vid.same_as(yy.args[0].vid)
@@ -255,8 +251,11 @@ def test_type_args():
     x = relay.var("x", shape=(10, 10))
     y = relay.var("y", shape=(1, 10))
     z = relay.add(x, y)
-    ty_z = infer_expr(z)
-    ty_args = ty_z.type_args
+
+    # InferTypeLocal does not support populating the type_args field
+    mod = infer_mod(IRModule.from_expr(z))
+    mod = infer_mod(mod, annotate_spans=False)
+    ty_args = mod["main"].body.type_args
     assert len(ty_args) == 2
     assert ty_args[0].dtype == "float32"
     assert ty_args[1].dtype == "float32"

From faeed08aed8ea983abe8e05494c96bdf0494e74c Mon Sep 17 00:00:00 2001
From: Andrew Luo <andrew.zhao.luo@gmail.com>
Date: Wed, 22 Dec 2021 13:29:15 -0800
Subject: [PATCH 19/19] jostle