From ee8420aefadae201d9100e7fed649a68b901cb39 Mon Sep 17 00:00:00 2001
From: Wei Luning <weiluning@huawei.com>
Date: Tue, 9 Jun 2020 21:03:29 +0800
Subject: [PATCH] Make assign operator to use the signature.

---
 mindspore/ccsrc/operator/ops.h                |  3 ++-
 mindspore/ccsrc/operator/ops_extends.cc       |  4 ++--
 .../ccsrc/pipeline/parse/function_block.cc    |  7 ++-----
 tests/ut/python/ops/test_signature.py         | 21 ++++++++++++++++++-
 4 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/mindspore/ccsrc/operator/ops.h b/mindspore/ccsrc/operator/ops.h
index 00484e1e1b..70100b85d6 100755
--- a/mindspore/ccsrc/operator/ops.h
+++ b/mindspore/ccsrc/operator/ops.h
@@ -27,7 +27,8 @@ namespace mindspore {
 // namespace to support primitive operators
 namespace prim {
 ValuePtr GetPythonOps(const std::string &op_name,
-                      const std::string &module_name = "mindspore._extends.parse.standard_method");
+                      const std::string &module_name = "mindspore._extends.parse.standard_method",
+                      bool use_signature = false);
 
 // Arithmetic
 extern const PrimitivePtr kPrimScalarAdd;
diff --git a/mindspore/ccsrc/operator/ops_extends.cc b/mindspore/ccsrc/operator/ops_extends.cc
index 6a192eca10..d415b45adf 100755
--- a/mindspore/ccsrc/operator/ops_extends.cc
+++ b/mindspore/ccsrc/operator/ops_extends.cc
@@ -23,10 +23,10 @@
 namespace mindspore {
 // namespace to support primitive operators
 namespace prim {
-ValuePtr GetPythonOps(const std::string &op_name, const std::string &module_name) {
+ValuePtr GetPythonOps(const std::string &op_name, const std::string &module_name, bool use_signature) {
   py::object obj = parse::python_adapter::GetPyFn(module_name, op_name);
   ValuePtr node = nullptr;
-  bool succ = parse::ConvertData(obj, &node);
+  bool succ = parse::ConvertData(obj, &node, use_signature);
   if (!succ) {
     MS_LOG(EXCEPTION) << "get Python op " << op_name << " from " << module_name << " fail";
   }
diff --git a/mindspore/ccsrc/pipeline/parse/function_block.cc b/mindspore/ccsrc/pipeline/parse/function_block.cc
index 2861da3116..fbeeba94a1 100644
--- a/mindspore/ccsrc/pipeline/parse/function_block.cc
+++ b/mindspore/ccsrc/pipeline/parse/function_block.cc
@@ -322,12 +322,10 @@ void FunctionBlock::InsertDependItemsBeforeReturn() {
 
   ValueNodePtr make_tuple_op = NewValueNode(prim::kPrimMakeTuple);
   ValueNodePtr depend_op = NewValueNode(prim::kPrimDepend);
-  ValueNodePtr get_ref_origin_op = NewValueNode(prim::kPrimGetRefOrigin);
   ValueNodePtr stop_gradient_op = NewValueNode(prim::kPrimStopGradient);
   const std::string primitive_name("assign");
   const std::string module_name("mindspore.ops.functional");
-  ValueNodePtr assign_op = NewValueNode(prim::GetPythonOps(primitive_name, module_name));
-
+  ValueNodePtr assign_op = NewValueNode(prim::GetPythonOps(primitive_name, module_name, true));
   if (state_assign_.size() == 0 && auto_depends_.size() == 0) {
     return;
   }
@@ -336,8 +334,7 @@ void FunctionBlock::InsertDependItemsBeforeReturn() {
   vec_states.emplace_back(make_tuple_op);
   for (auto &item : state_assign_) {
     auto source = ReadVariable(item.second);
-    auto origin = func_graph()->NewCNode({get_ref_origin_op, item.first});
-    auto assign = func_graph()->NewCNode({assign_op, origin, source});
+    auto assign = func_graph()->NewCNode({assign_op, item.first, source});
     MS_LOG(INFO) << "SetState read " << item.first->ToString() << ", " << item.second;
     vec_states.emplace_back(assign);
   }
diff --git a/tests/ut/python/ops/test_signature.py b/tests/ut/python/ops/test_signature.py
index e6447be8f3..6de00330af 100644
--- a/tests/ut/python/ops/test_signature.py
+++ b/tests/ut/python/ops/test_signature.py
@@ -47,7 +47,7 @@ class Net(nn.Cell):
 
 
 def test_assign_through_cell():
-    context.set_context(mode=context.GRAPH_MODE, save_graphs=True)
+    context.set_context(mode=context.GRAPH_MODE)
     net = Net()
     net.to_float(ms.float16)
     net.add_flags_recursive(fp16=False)
@@ -57,6 +57,25 @@ def test_assign_through_cell():
         net(None)
 
 
+class AssignOp(nn.Cell):
+    def __init__(self):
+        super(AssignOp, self).__init__()
+        self.b = Parameter(initializer('ones', [5]), name='b')
+
+
+    def construct(self, w):
+        self.b = w
+        return w
+
+
+def test_assign_by_operator():
+    context.set_context(mode=context.GRAPH_MODE)
+    net = AssignOp()
+    net.to_float(ms.float16)
+    input_data = Tensor(np.ones([5]).astype(np.float32))
+    net(input_data)
+
+
 class NetScatterNdUpdate(nn.Cell):
     def __init__(self):
         super(NetScatterNdUpdate, self).__init__()