From 4ad44385321bd78199c8a0fe0c47d422d401faf2 Mon Sep 17 00:00:00 2001
From: Marisa Kirisame <lolisa@marisa.moe>
Date: Thu, 10 Sep 2020 03:28:30 +0000
Subject: [PATCH] upstream

---
 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 0ecbfea8c905..7a1941b161a1 100644
--- a/src/relay/transforms/fold_constant.cc
+++ b/src/relay/transforms/fold_constant.cc
@@ -104,7 +104,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"};