[DoubleGrad #4] Bug Fixes to Double Grad Node Generation (#41121)

* [Refactor] refactored eager_gen.py PR #2

* [DoubleGrad PR #1] Decoupled code generation logics for Dygraph ForwardFunctions and GradNodes

* Fixed minor issue

* Adjusted logics of GenerateNodeCreationCodes and GenerateForwardDefinition

* Fixed issues

* Supported higher-order grad node generation

* [DoubleGrad PR #4] Supported higher-order GradNode generation

* [DoubleGrad #4] Bug Fixes to Double Grad Node Generation

* Fixed yaml typo

* Fixed yaml typo

* fixed minor issues

* Fixed minor issue

paddle/fluid/eager/auto_code_generator/final_state_generator/codegen_utils.py

@@ -21,7 +21,8 @@ import os
 ########################
 ### Global Variables ###
 ########################
-ops_to_fill_zero_for_empty_grads = set(["split", "rnn"])
+ops_to_fill_zero_for_empty_grads = set(
+    ["split_grad", "rnn_grad", "matmul_double_grad"])
 
 # For API dispatch used at python-level
 # { op_name : [arg_name, ...] }
@@ -176,6 +177,11 @@ def TransformGradVarNameForDoubleGradGeneration(string):
     return string
 
 
+def GetIndent(num):
+    tab = "   "
+    return "".join([tab for i in range(num)])
+
+
 ######################
 ###  Yaml Parsers  ###
 ######################

paddle/fluid/eager/grad_node_info.h

@@ -14,6 +14,8 @@
 
 #pragma once
 
+#include <memory>
+
 #include "paddle/fluid/eager/eager_tensor.h"
 #include "paddle/fluid/eager/hooks.h"
 #include "paddle/phi/api/all.h"

python/paddle/nn/functional/activation.py

@@ -23,6 +23,7 @@ from ...tensor.math import multiply
 import warnings
 from ...fluid.layer_helper import LayerHelper
 from ...fluid.framework import convert_np_dtype_to_dtype_
+from ...fluid.framework import _in_legacy_dygraph, in_dygraph_mode
 from ...fluid.data_feeder import check_variable_and_dtype, check_dtype
 import paddle
 from paddle import _C_ops, in_dynamic_mode
@@ -560,9 +561,10 @@ def relu(x, name=None):
             out = F.relu(x) # [0., 0., 1.]
     """
 
-    if in_dynamic_mode():
+    if in_dygraph_mode():
+        return _C_ops.final_state_relu(x)
+    if _in_legacy_dygraph():
         return _C_ops.relu(x)
-
     check_variable_and_dtype(x, 'x', ['float16', 'float32', 'float64'], 'relu')
     helper = LayerHelper('relu', **locals())
     out = helper.create_variable_for_type_inference(x.dtype)

python/paddle/utils/code_gen/backward.yaml

@@ -377,15 +377,15 @@
     data_type : x
 
 - backward_api : matmul_double_grad
-  forward : matmul_grad (Tensor x, Tensor y, Tensor out_grad, bool transpose_x, bool transpose_y) -> Tensor(dx), Tensor(dy)
-  args : (Tensor x, Tensor y, Tensor out_grad, Tensor dx_grad, Tensor dy_grad, bool transpose_x, bool transpose_y)
-  output : Tensor(d2x), Tensor(d2y), Tensor(dout_grad)
+  forward : matmul_grad (Tensor x, Tensor y, Tensor grad_out, bool transpose_x=false, bool transpose_y=false) -> Tensor(grad_x), Tensor(grad_y)
+  args : (Tensor x, Tensor y, Tensor grad_out, Tensor grad_x_grad, Tensor grad_y_grad, bool transpose_x=false, bool transpose_y=false)
+  output : Tensor(x_grad), Tensor(y_grad), Tensor(grad_out_grad)
   infer_meta :
     func : GeneralTernaryGradInferMeta
-    param : [x, y, out_grad]
+    param : [x, y, grad_out]
   kernel :
     func : matmul_double_grad
-  optional : dx_grad, dy_grad
+  optional : grad_x_grad, grad_y_grad
 
 - backward_api : matmul_grad
   forward : matmul (Tensor x, Tensor y, bool transpose_x=false, bool transpose_y=false) -> Tensor(out)
@@ -396,6 +396,7 @@
     param : [x, y]
   kernel :
     func : matmul_grad
+  backward : matmul_double_grad
 
 - backward_api : matrix_power_grad
   forward : matrix_power (Tensor x, int n) -> Tensor(out)