From be07d46a87a1bb1d2ff8190e7d47cd9cac7c7421 Mon Sep 17 00:00:00 2001
From: linqingke <linqingke@huawei.com>
Date: Thu, 25 Mar 2021 16:19:05 +0800
Subject: [PATCH] revert Remove uselesee updatestate pattern and fix loss and J
 order

---
 .../optimizer/auto_monad_eliminate.cc         | 25 ++++++++++++++++---
 mindspore/nn/wrap/cell_wrapper.py             |  1 -
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/mindspore/ccsrc/frontend/optimizer/auto_monad_eliminate.cc b/mindspore/ccsrc/frontend/optimizer/auto_monad_eliminate.cc
index 1102d3cca3..be31d880e9 100644
--- a/mindspore/ccsrc/frontend/optimizer/auto_monad_eliminate.cc
+++ b/mindspore/ccsrc/frontend/optimizer/auto_monad_eliminate.cc
@@ -107,6 +107,20 @@ std::vector<std::vector<size_t>> SplitGroup(const std::vector<AnfNodePtr> &topos
 // a = Load(para1, u1)
 // ...
 // b = Load(para1, u2)
+// u3 = UpdateState(u2, b)
+//==>
+// delete the UpdateState
+void DeleteLoadUserUpdateState(const FuncGraphManagerPtr &manager, const AnfNodePtr &load_user,
+                               const AnfNodePtr &load) {
+  const auto &load_cnode = load->cast<CNodePtr>();
+  const auto &u = load_cnode->input(2);
+  manager->Replace(load_user, u);
+}
+
+// Pattern2======================================
+// a = Load(para1, u1)
+// ...
+// b = Load(para1, u2)
 // t = make_tuple(x, b)
 // u3 = UpdateState(u2, t)
 //==>
@@ -127,7 +141,7 @@ void DeleteLoadUserMakeTuple(const FuncGraphManagerPtr &manager, const CNodePtr
   manager->Replace(make_tuple, other_input);
 }
 
-// Pattern2======================================
+// Pattern3======================================
 // a = Load(para1, u1)
 // ...
 // b = Load(para1, u2)
@@ -153,6 +167,11 @@ void ReplaceLoadUserMakeTuple(const FuncGraphManagerPtr &manager, const FuncGrap
 void ReplaceLoadUser(const FuncGraphManagerPtr &manager, const FuncGraphPtr &fg, const AnfNodePtr &load) {
   auto load_users = manager->node_users()[load];
   for (const auto &load_user : load_users) {
+    // Pattern1
+    if (IsPrimitiveCNode(load_user.first, prim::kPrimUpdateState)) {
+      DeleteLoadUserUpdateState(manager, load_user.first, load);
+      continue;
+    }
     if (IsPrimitiveCNode(load_user.first, prim::kPrimMakeTuple)) {
       const auto &make_tuple = load_user.first->cast<CNodePtr>();
       auto &maketuple_users = manager->node_users()[make_tuple];
@@ -161,12 +180,12 @@ void ReplaceLoadUser(const FuncGraphManagerPtr &manager, const FuncGraphPtr &fg,
       if (!maketuple_as_input_of_update) {
         continue;
       }
-      // Pattern1
+      // Pattern2
       if (make_tuple->size() == 3) {
         DeleteLoadUserMakeTuple(manager, make_tuple, load);
         continue;
       }
-      // Pattern2
+      // Pattern3
       if (make_tuple->size() > 3) {
         ReplaceLoadUserMakeTuple(manager, fg, make_tuple, load);
       }
diff --git a/mindspore/nn/wrap/cell_wrapper.py b/mindspore/nn/wrap/cell_wrapper.py
index b143a736bc..003e257149 100644
--- a/mindspore/nn/wrap/cell_wrapper.py
+++ b/mindspore/nn/wrap/cell_wrapper.py
@@ -352,7 +352,6 @@ class TrainOneStepCell(Cell):
         weights = self.weights
         loss = self.network(*inputs)
         sens = P.Fill()(P.DType()(loss), P.Shape()(loss), self.sens)
-        sens = F.depend(sens, loss)
         grads = self.grad(self.network, weights)(*inputs, sens)
         grads = self.grad_reducer(grads)
         loss = F.depend(loss, self.optimizer(grads))