From ea4c9fcf1e408fc3a0fe2d9fb58bd6a004fb0609 Mon Sep 17 00:00:00 2001
From: Aurelius84 <zhangliujie@baidu.com>
Date: Fri, 8 Apr 2022 16:08:52 +0800
Subject: [PATCH] [Eager]Fix segment_pool/allclose/isclose/scale API bug
 (#41506)

* [Eager]Fix segment_pool/allclose/isclose/scale API bug

* fix kernel register problem
---
 paddle/fluid/operators/cast_op.cu          | 22 +++++++++++-----------
 python/paddle/incubate/tensor/math.py      |  2 +-
 python/paddle/tensor/logic.py              | 14 ++++++++++++--
 python/paddle/utils/code_gen/backward.yaml |  4 +++-
 4 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/paddle/fluid/operators/cast_op.cu b/paddle/fluid/operators/cast_op.cu
index eb51215790bbc..0afe09ec028e3 100644
--- a/paddle/fluid/operators/cast_op.cu
+++ b/paddle/fluid/operators/cast_op.cu
@@ -19,15 +19,15 @@ namespace ops = paddle::operators;
 namespace plat = paddle::platform;
 
 using CUDA = paddle::platform::CUDADeviceContext;
-#define REGISTER_CAST_CUDA_BASE(op_name, ...)                             \
-  REGISTER_OP_CUDA_KERNEL(                                                \
-      op_name, ops::CastOpKernel<CUDA, float>,                            \
-      ops::CastOpKernel<CUDA, double>, ops::CastOpKernel<CUDA, int>,      \
-      ops::CastOpKernel<CUDA, int64_t>, ops::CastOpKernel<CUDA, int16_t>, \
-      ops::CastOpKernel<CUDA, bool>, ops::CastOpKernel<CUDA, uint8_t>,    \
-      ops::CastOpKernel<CUDA, plat::float16>,                             \
-      ops::CastOpKernel<CUDA, plat::complex<float>>,                      \
-      ops::CastOpKernel<CUDA, plat::complex<double>>, ##__VA_ARGS__);
-
 // See [ why register transfer_dtype_op alias with cast_op? ] in cast_op.cc
-REGISTER_CAST_CUDA_BASE(transfer_dtype, ops::CastOpKernel<CUDA, plat::bfloat16>)
+REGISTER_OP_CUDA_KERNEL(transfer_dtype, ops::CastOpKernel<CUDA, float>,
+                        ops::CastOpKernel<CUDA, double>,
+                        ops::CastOpKernel<CUDA, int>,
+                        ops::CastOpKernel<CUDA, int64_t>,
+                        ops::CastOpKernel<CUDA, int16_t>,
+                        ops::CastOpKernel<CUDA, bool>,
+                        ops::CastOpKernel<CUDA, uint8_t>,
+                        ops::CastOpKernel<CUDA, plat::float16>,
+                        ops::CastOpKernel<CUDA, plat::complex<float>>,
+                        ops::CastOpKernel<CUDA, plat::complex<double>>,
+                        ops::CastOpKernel<CUDA, plat::bfloat16>);
diff --git a/python/paddle/incubate/tensor/math.py b/python/paddle/incubate/tensor/math.py
index b36aaef9acf36..6f88e94b7ee63 100644
--- a/python/paddle/incubate/tensor/math.py
+++ b/python/paddle/incubate/tensor/math.py
@@ -222,7 +222,7 @@ def segment_max(data, segment_ids, name=None):
     """
 
     if in_dygraph_mode():
-        out, tmp = _C_ops.final_state_segment_pool(data, segment_ids, "MAX")[0]
+        out, tmp = _C_ops.final_state_segment_pool(data, segment_ids, "MAX")
         return out
 
     if _non_static_mode():
diff --git a/python/paddle/tensor/logic.py b/python/paddle/tensor/logic.py
index 03d0f42d8417b..972a2f7296877 100755
--- a/python/paddle/tensor/logic.py
+++ b/python/paddle/tensor/logic.py
@@ -123,7 +123,12 @@ def allclose(x, y, rtol=1e-05, atol=1e-08, equal_nan=False, name=None):
     """
 
     if in_dygraph_mode():
-        return _C_ops.final_state_allclose(x, y, rtol, atol, equal_nan)
+        # NOTE(dev): Pass tol as Tensor to fix precision loss problem, because
+        # C++ backend will cast it into float32 if passing float from python.
+        as_tensor = lambda x: paddle.to_tensor([x], dtype='float64', place='cpu')
+        return _C_ops.final_state_allclose(x, y,
+                                           as_tensor(rtol),
+                                           as_tensor(atol), equal_nan)
     if _in_legacy_dygraph():
         return _C_ops.allclose(x, y, 'rtol',
                                str(rtol), 'atol',
@@ -685,7 +690,12 @@ def isclose(x, y, rtol=1e-05, atol=1e-08, equal_nan=False, name=None):
     """
 
     if in_dygraph_mode():
-        return _C_ops.final_state_isclose(x, y, rtol, atol, equal_nan)
+        # NOTE(dev): Pass tol as Tensor to fix precision loss problem, because
+        # C++ backend will cast it into float32 if passing float from python.
+        as_tensor = lambda x: paddle.to_tensor([x], dtype='float64', place='cpu')
+        return _C_ops.final_state_isclose(x, y,
+                                          as_tensor(rtol),
+                                          as_tensor(atol), equal_nan)
     if _in_legacy_dygraph():
         return _C_ops.isclose(x, y, 'rtol',
                               str(rtol), 'atol',
diff --git a/python/paddle/utils/code_gen/backward.yaml b/python/paddle/utils/code_gen/backward.yaml
index 4cb411634a0ad..32a54f3dffd0d 100644
--- a/python/paddle/utils/code_gen/backward.yaml
+++ b/python/paddle/utils/code_gen/backward.yaml
@@ -1096,7 +1096,7 @@
   forward : scale (Tensor x, Scalar scale, float bias, bool bias_after_scale) -> Tensor(out)
   args : (Tensor out_grad, Scalar scale=1.0, float bias=0.0, bool bias_after_scale=true)
   output : Tensor(x_grad)
-  invoke : scale(out_grad, scale, bias, bias_after_scale)
+  invoke : scale(out_grad, scale, 0.0, bias_after_scale)
 
 - backward_api : scatter_grad
   forward : scatter (Tensor x, Tensor index, Tensor updates, bool overwrite) -> Tensor(out)
@@ -1129,6 +1129,8 @@
     param : [x]
   kernel :
     func : segment_pool_grad
+    data_type : x
+  optional : summed_ids
 
 - backward_api : selu_grad
   forward : selu (Tensor x, float scale, float alpha) -> Tensor(out)