From fd591ecb457d0f7f76ef6ddaa6c2ef02248bdb5f Mon Sep 17 00:00:00 2001
From: Aurelius84 <zhangliujie@baidu.com>
Date: Mon, 4 Apr 2022 06:04:33 +0800
Subject: [PATCH] [Eager]Polish enable/disable_legacy_dygraph logic (#41364)

* [Eager]Polish enable/disable_legacy_dygraph logic

* merge yunfei PR

* merge other pr
---
 python/paddle/fluid/framework.py     | 53 ++++++++++++++--------------
 python/paddle/tensor/manipulation.py | 35 ++++++++----------
 2 files changed, 40 insertions(+), 48 deletions(-)

diff --git a/python/paddle/fluid/framework.py b/python/paddle/fluid/framework.py
index 20c441f364..a329610eea 100644
--- a/python/paddle/fluid/framework.py
+++ b/python/paddle/fluid/framework.py
@@ -115,14 +115,38 @@ def _update_monkey_methods(is_eager):
     from .dygraph.varbase_patch_methods import monkey_patch_varbase
     from .dygraph import monkey_patch_math_varbase
 
+    global _already_patch_eager_tensor
+    global _already_patch_varbase
+
     assert isinstance(is_eager, bool)
+    # switch into eager mode
     if is_eager:
         _C_ops.switch_to_eager_ops()
+        if not _already_patch_eager_tensor:
+            monkey_patch_varbase()
+            monkey_patch_math_varbase()
+
+            _already_patch_eager_tensor = True
+    # switch back into legacy mode
     else:
         _C_ops.switch_to_core_ops()
+        if not _already_patch_varbase:
+            monkey_patch_varbase()
+            monkey_patch_math_varbase()
+
+            _already_patch_varbase = True
 
-    monkey_patch_varbase()
-    monkey_patch_math_varbase()
+    # switch Paddle.Tensor bind type
+    _switch_tensor_bind_type(is_eager)
+
+
+def _switch_tensor_bind_type(is_eager):
+    import paddle
+    if is_eager:
+        paddle.Tensor = core.eager.Tensor
+    else:
+        paddle.Tensor = core.VarBase
+    paddle.Tensor.__qualname__ = 'Tensor'
 
 
 def _enable_legacy_dygraph():
@@ -183,35 +207,10 @@ def _non_static_mode():
 @signature_safe_contextmanager
 def _test_eager_guard(place=None):
     _disable_legacy_dygraph()
-    from paddle import _C_ops
-    _C_ops.switch_to_eager_ops()
-    global _already_patch_eager_tensor
-    global _already_patch_varbase
-    from .dygraph.varbase_patch_methods import monkey_patch_varbase
-    from .dygraph import monkey_patch_math_varbase
-    if not _already_patch_eager_tensor:
-        monkey_patch_varbase()
-        monkey_patch_math_varbase()
-
-        # Ugly setting
-        from paddle.tensor.manipulation import fill_, zero_, fill_diagonal_, fill_diagonal_tensor_, tolist
-        setattr(core.eager.Tensor, 'fill_', fill_)
-        setattr(core.eager.Tensor, 'zero_', zero_)
-        setattr(core.eager.Tensor, 'fill_diagonal_', fill_diagonal_)
-        setattr(core.eager.Tensor, 'fill_diagonal_tensor_',
-                fill_diagonal_tensor_)
-        setattr(core.eager.Tensor, 'tolist', tolist)
-
-        _already_patch_eager_tensor = True
     try:
         yield
     finally:
         _enable_legacy_dygraph()
-        if not _already_patch_varbase:
-            monkey_patch_varbase()
-            monkey_patch_math_varbase()
-            _already_patch_varbase = True
-        _C_ops.switch_to_core_ops()
 
 
 global_ipu_index = None
diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py
index ca807c286a..f6bbadf987 100755
--- a/python/paddle/tensor/manipulation.py
+++ b/python/paddle/tensor/manipulation.py
@@ -76,9 +76,6 @@ def fill_(x, value):
                             float(value), "value_int", int(value))
 
 
-setattr(core.VarBase, 'fill_', fill_)
-
-
 @dygraph_only
 def zero_(x):
     """
@@ -107,9 +104,6 @@ def zero_(x):
     return _C_ops.fill_any_(x, "value_float", 0., "value_int", int(0))
 
 
-setattr(core.VarBase, 'zero_', zero_)
-
-
 @dygraph_only
 def fill_diagonal_(x, value, offset=0, wrap=False, name=None):
     """
@@ -156,9 +150,6 @@ def fill_diagonal_(x, value, offset=0, wrap=False, name=None):
                                  True)
 
 
-setattr(core.VarBase, 'fill_diagonal_', fill_diagonal_)
-
-
 def _fill_diagonal_tensor_impl(x, y, offset=0, dim1=0, dim2=1, inplace=False):
     inshape = x.shape
     assert dim1 < len(inshape) and dim1 >= -len(inshape), (
@@ -226,9 +217,6 @@ def fill_diagonal_tensor_(x, y, offset=0, dim1=0, dim2=1, name=None):
         x, y, offset=offset, dim1=dim1, dim2=dim2, inplace=True)
 
 
-setattr(core.VarBase, 'fill_diagonal_tensor_', fill_diagonal_tensor_)
-
-
 def fill_diagonal_tensor(x, y, offset=0, dim1=0, dim2=1, name=None):
     """
     This function fill the source Tensor y into the x Tensor's diagonal.
@@ -262,12 +250,6 @@ def fill_diagonal_tensor(x, y, offset=0, dim1=0, dim2=1, name=None):
         x, y, offset=offset, dim1=dim1, dim2=dim2, inplace=False)
 
 
-setattr(core.VarBase, 'fill_diagonal_tensor', fill_diagonal_tensor)
-
-if _in_eager_without_dygraph_check():
-    setattr(core.eager.Tensor, 'fill_diagonal_tensor', fill_diagonal_tensor)
-
-
 @dygraph_only
 def tolist(x):
     """
@@ -301,9 +283,6 @@ def tolist(x):
     return x.numpy().tolist()
 
 
-setattr(core.VarBase, 'tolist', tolist)
-
-
 def concat(x, axis=0, name=None):
     """
 
@@ -2961,3 +2940,17 @@ def put_along_axis_(arr, indices, values, axis, reduce='assign'):
     values = paddle.broadcast_to(values, indices.shape)
     return _C_ops.put_along_axis_(arr, indices, values, "Axis", axis, "Reduce",
                                   reduce)
+
+
+# TODO(dev): We need avoid implementing it by this way.
+__METHODS = {
+    'fill_': fill_,
+    'zero_': zero_,
+    'fill_diagonal_': fill_diagonal_,
+    'fill_diagonal_tensor_': fill_diagonal_tensor_,
+    "fill_diagonal_tensor": fill_diagonal_tensor,
+    'tolist': tolist
+}
+for name, func in __METHODS.items():
+    setattr(core.VarBase, name, func)
+    setattr(core.eager.Tensor, name, func)
-- 
GitLab