From 9b0e37a42de0f4cb0efcaa907bfa54d61f1c3faf Mon Sep 17 00:00:00 2001
From: =?utf8?q?=E9=9B=BE=E9=9B=A8=E9=AD=94=E7=90=86=E6=B2=99?=
 <lolisa@marisa.moe>
Date: Fri, 11 Sep 2020 17:00:50 -0700
Subject: [PATCH] upstream (#6436)

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

diff --git a/src/relay/transforms/fold_constant.cc b/src/relay/transforms/fold_constant.cc
index 0a1f173..a3f2f69 100644
--- a/src/relay/transforms/fold_constant.cc
+++ b/src/relay/transforms/fold_constant.cc
@@ -105,7 +105,23 @@ class ConstantFolder : public ExprMutator {
     }
   }
 
+  bool inside_primitive = false;
+  Expr VisitExpr_(const FunctionNode* op) final {
+    if (op->HasNonzeroAttr(attr::kPrimitive)) {
+      CHECK_EQ(inside_primitive, false);
+      inside_primitive = true;
+      auto ret = ExprMutator::VisitExpr_(op);
+      inside_primitive = false;
+      return ret;
+    } else {
+      return ExprMutator::VisitExpr_(op);
+    }
+  }
+
   Expr VisitExpr_(const CallNode* call) final {
+    if (inside_primitive) {
+      return GetRef<Expr>(call);
+    }
     static auto op_stateful = Op::GetAttrMap<TOpIsStateful>("TOpIsStateful");
 
     std::unordered_set<std::string> skip_list{"zeros_like", "ones_like", "full_like", "full"};
-- 
2.7.4