fix prim abs_grad nan bug (#57687)

cyber-pioneer · web-flow · commit 81d15b488995 · 2023-10-11T14:00:44.000+08:00
diff --git a/paddle/fluid/prim/api/api.yaml b/paddle/fluid/prim/api/api.yaml
@@ -48,3 +48,4 @@
 - reshape
 - erf
 - tanh
+- sign
diff --git a/paddle/fluid/prim/api/composite_backward/composite_backward_api.h b/paddle/fluid/prim/api/composite_backward/composite_backward_api.h
@@ -585,9 +585,8 @@ void sigmoid_grad(const Tensor& out, const Tensor& out_grad, Tensor* x_grad) {
 template <typename T>
 void abs_grad(const Tensor& x, const Tensor& out_grad, Tensor* x_grad) {
   if (x_grad) {
-    auto abs_tmp = abs<T>(x);
-    auto divide_tmp = divide<T>(x, abs_tmp);
-    set_output<T>(out_grad * divide_tmp, x_grad);
+    auto sign_tmp = sign<T>(x);
+    set_output<T>(out_grad * sign_tmp, x_grad);
   }
 }
 
diff --git a/paddle/fluid/primitive/codegen/gen.py b/paddle/fluid/primitive/codegen/gen.py
@@ -294,6 +294,7 @@
     'diag_grad',
     'trace_grad',
     'flip',
+    'sign',
 ]
 
 
diff --git a/paddle/fluid/primitive/primitive.yaml b/paddle/fluid/primitive/primitive.yaml
@@ -50,3 +50,4 @@
 - tanh
 - full
 - cast
+- sign
diff --git a/paddle/fluid/primitive/rule/vjp/manual/manual_vjp.cc b/paddle/fluid/primitive/rule/vjp/manual/manual_vjp.cc
@@ -55,6 +55,7 @@ std::vector<std::vector<paddle::Tensor>> reshape_vjp(
   if (paddle::prim::StaticCompositeContext::Instance().IsBwdPrimEnabled() &&
       !need_skip) {
     FLAGS_tensor_operants_mode = "static";
+    VLOG(4) << "Call PIR Decomposed backward op reshape_grad";
     paddle::Tensor* x_grad = !stop_gradients[0][0] ? &vjp_res[0][0] : nullptr;
 
     details::reshape_grad<LazyTensor>(xshape, out_grad, x_grad);

-Original file line number
+Diff line change
 - reshape
 - erf
 - tanh
 +- sign
Original file line number	Diff line number	Diff line change
`@@ -585,9 +585,8 @@ void sigmoid_grad(const Tensor& out, const Tensor& out_grad, Tensor* x_grad) {`
`585`	`585`	`template <typename T>`
`586`	`586`	`void abs_grad(const Tensor& x, const Tensor& out_grad, Tensor* x_grad) {`
`587`	`587`	`if (x_grad) {`
`588`		`- auto abs_tmp = abs<T>(x);`
`589`		`- auto divide_tmp = divide<T>(x, abs_tmp);`
`590`		`- set_output<T>(out_grad * divide_tmp, x_grad);`
	`588`	`+ auto sign_tmp = sign<T>(x);`
	`589`	`+ set_output<T>(out_grad * sign_tmp, x_grad);`
`591`	`590`	`}`
`592`	`591`	`}`
`593`	`592`
Original file line number	Diff line number	Diff line change
`@@ -294,6 +294,7 @@`
`294`	`294`	`'diag_grad',`
`295`	`295`	`'trace_grad',`
`296`	`296`	`'flip',`
	`297`	`+ 'sign',`
`297`	`298`	`]`
`298`	`299`
`299`	`300`