From 50714d5cc41d121b9bb979023bc58eabc2a3a49a Mon Sep 17 00:00:00 2001
From: Aurelius84 <zhangliujie@baidu.com>
Date: Sat, 2 Apr 2022 21:07:11 +0800
Subject: [PATCH] [Eager]Fix eager no take effect problem (#41291)

* [Eager]Fix eager no take effect problem

* add element_wise and fix greater_than
---
 paddle/fluid/pybind/eager_method.cc                  | 11 +++++++++++
 python/paddle/__init__.py                            |  5 ++++-
 python/paddle/fluid/tests/unittests/test_cross_op.py |  4 ++--
 python/paddle/tensor/linalg.py                       |  4 ++++
 python/paddle/tensor/logic.py                        |  3 ++-
 python/paddle/utils/code_gen/api.yaml                |  8 ++++----
 6 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/paddle/fluid/pybind/eager_method.cc b/paddle/fluid/pybind/eager_method.cc
index 37ace14d145..d9face124bd 100644
--- a/paddle/fluid/pybind/eager_method.cc
+++ b/paddle/fluid/pybind/eager_method.cc
@@ -1279,6 +1279,15 @@ static PyObject* tensor__inplace_version(TensorObject* self, PyObject* args,
   EAGER_CATCH_AND_THROW_RETURN_NULL
 }
 
+static PyObject* tensor_method_element_size(TensorObject* self, PyObject* args,
+                                            PyObject* kwargs) {
+  EAGER_TRY
+  uint32_t element_size = framework::DataTypeSize(self->tensor.dtype());
+
+  return ToPyObject(element_size);
+  EAGER_CATCH_AND_THROW_RETURN_NULL
+}
+
 static PyObject* tensor__bump_inplace_version(TensorObject* self,
                                               PyObject* args,
                                               PyObject* kwargs) {
@@ -1417,6 +1426,8 @@ PyMethodDef variable_methods[] = {
      METH_VARARGS | METH_KEYWORDS, NULL},
     {"to_dense", (PyCFunction)(void (*)(void))tensor_method_to_dense,
      METH_VARARGS | METH_KEYWORDS, NULL},
+    {"element_size", (PyCFunction)(void (*)(void))tensor_method_element_size,
+     METH_VARARGS | METH_KEYWORDS, NULL},
     /***the method of sparse tensor****/
     {"_inplace_version", (PyCFunction)(void (*)(void))tensor__inplace_version,
      METH_VARARGS | METH_KEYWORDS, NULL},
diff --git a/python/paddle/__init__.py b/python/paddle/__init__.py
index bba9c226dc0..e532633b6eb 100755
--- a/python/paddle/__init__.py
+++ b/python/paddle/__init__.py
@@ -48,7 +48,10 @@ from .framework.dtype import bfloat16  # noqa: F401
 from .framework.dtype import bool  # noqa: F401
 from .framework.dtype import complex64  # noqa: F401
 from .framework.dtype import complex128  # noqa: F401
-from .framework import VarBase as Tensor  # noqa: F401
+if fluid.framework._in_eager_mode_:
+    Tensor = framework.core.eager.Tensor
+else:
+    from .framework import VarBase as Tensor  # noqa: F401
 Tensor.__qualname__ = 'Tensor'  # noqa: F401
 import paddle.compat  # noqa: F401
 import paddle.distributed  # noqa: F401
diff --git a/python/paddle/fluid/tests/unittests/test_cross_op.py b/python/paddle/fluid/tests/unittests/test_cross_op.py
index 6cba72213ff..8b884583646 100644
--- a/python/paddle/fluid/tests/unittests/test_cross_op.py
+++ b/python/paddle/fluid/tests/unittests/test_cross_op.py
@@ -48,10 +48,10 @@ class TestCrossOp(OpTest):
         self.outputs = {'Out': np.array(z_list).reshape(self.shape)}
 
     def test_check_output(self):
-        self.check_output(check_eager=False)
+        self.check_output(check_eager=True)
 
     def test_check_grad_normal(self):
-        self.check_grad(['X', 'Y'], 'Out', check_eager=False)
+        self.check_grad(['X', 'Y'], 'Out', check_eager=True)
 
 
 class TestCrossOpCase1(TestCrossOp):
diff --git a/python/paddle/tensor/linalg.py b/python/paddle/tensor/linalg.py
index 818ce2f5c67..8afab2e05f2 100644
--- a/python/paddle/tensor/linalg.py
+++ b/python/paddle/tensor/linalg.py
@@ -27,6 +27,9 @@ from paddle import _C_ops
 
 __all__ = []
 
+# Consistent with kDefaultDim from C++ Backend
+K_DEFAULT_DIM = 9
+
 
 def matmul(x, y, transpose_x=False, transpose_y=False, name=None):
     """
@@ -1157,6 +1160,7 @@ def cross(x, y, axis=None, name=None):
             #  [0. 0. 0.]]
     """
     if in_dygraph_mode():
+        axis = K_DEFAULT_DIM if axis is None else axis
         return _C_ops.final_state_cross(x, y, axis)
     else:
         if _in_legacy_dygraph():
diff --git a/python/paddle/tensor/logic.py b/python/paddle/tensor/logic.py
index e3ffd36d779..3896fa535ff 100755
--- a/python/paddle/tensor/logic.py
+++ b/python/paddle/tensor/logic.py
@@ -280,7 +280,8 @@ def greater_than(x, y, name=None):
             print(result1)  # result1 = [False False True]
     """
     if in_dygraph_mode():
-        return _C_ops.final_state_greater_than(x, y)
+        axis = -1  # default value
+        return _C_ops.final_state_greater_than(x, y, axis)
     else:
         if _in_legacy_dygraph():
             return _C_ops.greater_than(x, y)
diff --git a/python/paddle/utils/code_gen/api.yaml b/python/paddle/utils/code_gen/api.yaml
index ece46837c6d..b46accfb11b 100644
--- a/python/paddle/utils/code_gen/api.yaml
+++ b/python/paddle/utils/code_gen/api.yaml
@@ -610,21 +610,21 @@
     func : gelu
   backward : gelu_grad
 
-- api : greater
+- api : greater_equal
   args : (Tensor x, Tensor y, int axis = -1)
   output : Tensor
   infer_meta :
     func : CompareInferMeta
   kernel :
-    func : greater
+    func : greater_equal
 
-- api : greater_equal
+- api : greater_than
   args : (Tensor x, Tensor y, int axis = -1)
   output : Tensor
   infer_meta :
     func : CompareInferMeta
   kernel :
-    func : greater_equal
+    func : greater_than
 
 - api : gumbel_softmax
   args : (Tensor x, float temperature, bool hard, int axis)
-- 
GitLab