[Prim] Add reshape_double_grad and p_norm_grad for dygraph (#69239)

HydrogenSulfate · web-flow · commit 9b117be82562 · 2024-11-12T14:28:39.000+08:00
* support reshape_double_grad

* support composite p_norm_grad for dygraph

* update unitest

* update unitest

* update annotations of softmax_grad

* trigger CI

* fix indent

* update copyright time 2022-&gt;2024(test=document_fix)
diff --git a/paddle/fluid/eager/auto_code_generator/generator/eager_gen.py b/paddle/fluid/eager/auto_code_generator/generator/eager_gen.py
@@ -82,6 +82,7 @@
     "bmm_double_grad",
     "index_put_double_grad",
     "gather_nd_double_grad",
+    "reshape_double_grad",
 ]
 
 # white ops list whose kernel can automatically do type promotion.
@@ -2251,8 +2252,8 @@ def GenerateHigherOrderNodeCreationCode(self):
   if (!paddle::prim::PrimCommonUtils::IsEagerPrimEnabled() || need_skip) {{
     if (trace_backward) {{
        PADDLE_THROW(common::errors::Unavailable(
-       \"The Op {self.backward_api_name} doesn't have any grad\"
-       \"op. If you don't intend calculating higher order\"
+       \"The Op {self.backward_api_name} doesn't have any grad \"
+       \"op. If you don't intend calculating higher order \"
        \"derivatives, please set `create_graph`to False.\"));
     }}
   }}
@@ -2271,8 +2272,8 @@ def GenerateHigherOrderNodeCreationCode(self):
         elif not is_invoke_forward_api and not is_composite_grad_api:
             next_grad_node_creation_str = f"""  if (trace_backward) {{
     PADDLE_THROW(common::errors::Unavailable(
-    \"The Op {self.backward_api_name} doesn't have any grad\"
-    \"op. If you don't intend calculating higher order\"
+    \"The Op {self.backward_api_name} doesn't have any grad \"
+    \"op. If you don't intend calculating higher order \"
     \"derivatives, please set `create_graph`to False.\"));
   }}"""
 
diff --git a/paddle/fluid/prim/api/api.yaml b/paddle/fluid/prim/api/api.yaml
@@ -51,3 +51,5 @@
 - sign
 - sigmoid
 - index_put
+- isnan
+- isfinite
diff --git a/paddle/fluid/prim/api/composite_backward/composite_backward_api.h b/paddle/fluid/prim/api/composite_backward/composite_backward_api.h
@@ -114,21 +114,15 @@ void softmax_grad(const Tensor& out,
                   int axis,
                   Tensor* x_grad) {
   if (x_grad) {
+    if (axis < 0) {
+      axis += out.dims().size();
+    }
     if (out_grad.dims().size() > 0) {
-      if (axis >= 0) {
-        auto new_out_grad = out_grad * out;
-        auto tmp_x_grad = new_out_grad -
-                          out * sum<T>(new_out_grad, {axis}, out.dtype(), true);
-        set_output<T>(tmp_x_grad, x_grad);
-      } else {
-        auto new_out_grad = out_grad * out;
-        auto tmp_x_grad =
-            new_out_grad - out * sum<T>(new_out_grad,
-                                        {out.dims().size() + axis},
-                                        out.dtype(),
-                                        true);
-        set_output<T>(tmp_x_grad, x_grad);
-      }
+      // dx = dy * y - y * (dy*y).sum(axis)
+      auto new_out_grad = out_grad * out;
+      auto tmp_x_grad =
+          new_out_grad - out * sum<T>(new_out_grad, {axis}, out.dtype(), true);
+      set_output<T>(tmp_x_grad, x_grad);
     } else {
       set_output<T>(
           full<T>(common::vectorize(out_grad.dims()), 0.0, out_grad.dtype()),
@@ -1947,5 +1941,107 @@ void scatter_nd_add_grad(const Tensor& index,
   }
 }
 
+template <typename T>
+void p_norm_grad(const Tensor& x,
+                 /*output of forward was reserved for efficient backward*/
+                 const Tensor& out,
+                 const Tensor& out_grad,
+                 float porder,
+                 int axis,
+                 float epsilon,
+                 bool keepdim,
+                 bool asvector,
+                 Tensor* x_grad) {
+  if (x_grad) {
+    if (axis < 0) {
+      axis += x.dims().size();
+    }
+
+    Tensor x_grad_tmp;
+    if (porder == 0.0) {
+      // dx = 0
+      x_grad_tmp = full<T>(x.shape(), 0, x.dtype());
+    } else {
+      /* generic case formula:
+          dx = {
+            dy * y^(1-p) * |x|^(p-1) * sgn(x), if p != +-inf,
+            dy * sgn(x) * (x==y), if p == +-inf.
+          }
+      */
+      Tensor expand_out = out;
+      Tensor expand_out_grad = out_grad;
+      // firstly expand output_grad to same ndim with x for convenience
+      if (!keepdim) {
+        auto expand_shape = common::vectorize(out_grad.dims());
+        expand_shape.insert(expand_shape.begin() + axis, 1);
+        expand_out_grad = reshape<T>(out_grad, expand_shape);
+        expand_out = reshape<T>(out, expand_shape);
+      }
+
+      if (porder == 1.0) {
+        // dx = dy * sign(x)
+        auto x_sign = sign<T>(x);
+        x_grad_tmp = x_sign * expand_out_grad;
+      } else if (porder == 2.0) {
+        // dx = dy * (x / y)
+        x_grad_tmp = x / expand_out;
+        // fill zero to avoid division by zero
+        auto _zero_tensor =
+            full<T>(common::vectorize(x.dims()), 0.0, x.dtype());
+        auto finite_mask = isfinite<T>(x_grad_tmp);
+        x_grad_tmp = where<T>(finite_mask, x_grad_tmp, _zero_tensor);
+        x_grad_tmp = expand_out_grad * (x_grad_tmp);
+
+      } else if (porder == INFINITY || porder == -INFINITY) {
+        // dy * sgn(x) * (x==y), if p == +-inf.
+        auto x_abs = abs<T>(x);
+        auto mask =
+            cast<T>(bitwise_or<T>(equal<T>(x_abs, expand_out), isnan<T>(x_abs)),
+                    expand_out.dtype());
+        auto x_sign = sign<T>(x);
+        x_grad_tmp =
+            x_sign * ((expand_out_grad /
+                       sum<T>(mask, {axis}, expand_out_grad.dtype(), true)) *
+                      mask);
+
+      } else if (porder < 1.0) {
+        // dx = dy * y^(1-p) * |x|^(p-1) * sgn(x)
+        auto x_sign = sign<T>(x);
+        auto x_abs_pow = abs<T>(x);
+        x_abs_pow = x_abs_pow.pow(porder - 1);
+
+        auto x_scaled = x_sign * x_abs_pow;
+        x_grad_tmp = x_scaled * expand_out_grad * expand_out.pow(1 - porder);
+
+      } else if (porder < 2.0) {
+        // dx = dy * y^(1-p) * |x|^(p-1) * sgn(x)
+        auto x_sign = sign<T>(x);
+        auto x_abs_pow = abs<T>(x);
+        x_abs_pow = x_abs_pow.pow(porder - 1);
+
+        // auto scale_v = expand_out_grad / expand_out.pow(porder - 1);
+        // auto _zero_tensor =
+        //     full<T>(common::vectorize(x.dims()), 0.0, x.dtype());
+        // auto out_non_zero_mask = not_equal<T>(expand_out, _zero_tensor);
+        // scale_v = scale_v * cast<T>(out_non_zero_mask, scale_v.dtype());
+        // x_grad_tmp = x_sign * x_abs_pow * scale_v;
+
+        auto scale_v = expand_out_grad * expand_out.pow(1 - porder);
+        x_grad_tmp = x_sign * x_abs_pow * scale_v;
+
+      } else {
+        // dx = dy * y^(1-p) * |x|^(p-1) * sgn(x)
+        auto x_sign = sign<T>(x);
+        auto x_abs_pow = abs<T>(x);
+        x_abs_pow = x_abs_pow.pow(porder - 1);
+
+        auto x_scaled = x_sign * x_abs_pow;
+        x_grad_tmp = x_scaled * expand_out_grad * expand_out.pow(1 - porder);
+      }
+    }
+    set_output<T>(x_grad_tmp, x_grad);
+  }
+}
+
 }  // namespace prim
 }  // namespace paddle
diff --git a/paddle/fluid/prim/api/composite_backward/composite_double_backward_api.h b/paddle/fluid/prim/api/composite_backward/composite_double_backward_api.h
@@ -1061,5 +1061,16 @@ void gather_nd_double_grad(const Tensor& grad_out,
   }
 }
 
+template <typename T>
+void reshape_double_grad(const Tensor& grad_out,
+                         const Tensor& grad_x_grad,
+                         Tensor* grad_out_grad) {
+  if (grad_out_grad) {
+    // ddout = reshape(ddx, ddout.shape)
+    auto grad_out_grad_tmp = reshape<T>(grad_x_grad, grad_out.shape());
+    set_output<T>(grad_out_grad_tmp, grad_out_grad);
+  }
+}
+
 }  // namespace prim
 }  // namespace paddle
diff --git a/paddle/fluid/primitive/decomp_rule/decomp_vjp/details.h b/paddle/fluid/primitive/decomp_rule/decomp_vjp/details.h
@@ -1041,24 +1041,18 @@ void softmax_grad(const Tensor& out,
                   int axis,
                   Tensor* x_grad) {
   if (x_grad) {
+    if (axis < 0) {
+      axis += out.dims().size();
+    }
+
     if (out_grad.dims().size() > 0) {
-      if (axis >= 0) {
-        auto new_out_grad = out_grad * out;
-        auto tmp_x_grad = new_out_grad -
-                          out * sum<T>(new_out_grad, {axis}, out.dtype(), true);
-        set_output<T>(tmp_x_grad, x_grad);
-      } else {
-        auto new_out_grad = out_grad * out;
-        auto tmp_x_grad =
-            new_out_grad - out * sum<T>(new_out_grad,
-                                        {out.dims().size() + axis},
-                                        out.dtype(),
-                                        true);
-        set_output<T>(tmp_x_grad, x_grad);
-      }
+      auto new_out_grad = out_grad * out;
+      auto tmp_x_grad =
+          new_out_grad - out * sum<T>(new_out_grad, {axis}, out.dtype(), true);
+      set_output<T>(tmp_x_grad, x_grad);
     } else {
-      Tensor zeros = full_scalar<T>(0.0, out.dtype());
-      set_output<T>(out_grad * zeros, x_grad);
+      auto zeros = full<T>(common::vectorize(out.dims()), 0, out.dtype());
+      set_output<T>(zeros, x_grad);
     }
   }
 }
diff --git a/paddle/phi/ops/yaml/backward.yaml b/paddle/phi/ops/yaml/backward.yaml
@@ -2216,6 +2216,7 @@
     spmd_rule : PNormGradInferSpmd
   kernel :
     func : p_norm_grad
+  composite: p_norm_grad(x, out, out_grad, porder, axis, epsilon, keepdim, asvector, x_grad)
 
 - backward_op : pad3d_double_grad
   forward : pad3d_grad(Tensor x, Tensor grad_out, IntArray paddings, str mode="constant", float pad_value=0.0, str data_format="NCDHW") -> Tensor(grad_x)
@@ -2585,6 +2586,7 @@
   kernel :
     func : reshape_double_grad
   no_need_buffer : grad_out
+  composite: reshape_double_grad(grad_out, grad_x_grad, grad_out_grad)
   inplace : (grad_x_grad -> grad_out_grad)
 
 - backward_op : reshape_grad
diff --git a/python/paddle/nn/functional/norm.py b/python/paddle/nn/functional/norm.py
@@ -335,8 +335,8 @@ def layer_norm(
             size :math:`[*, normalized_shape[0], normalized_shape[1], ..., normalized_shape[-1]]`.
             If it is a single integer, this module will normalize over the last dimension
             which is expected to be of that specific size.
-        weight(Tensor, optional): The weight tensor of batch_norm. Default: None.
-        bias(Tensor, optional): The bias tensor of batch_norm. Default: None.
+        weight(Tensor, optional): The weight tensor of layer_norm. Default: None.
+        bias(Tensor, optional): The bias tensor of layer_norm. Default: None.
         epsilon(float, optional): The small value added to the variance to prevent
             division by zero. Default: 1e-05.
         name(str, optional): Name for the LayerNorm, default is None. For more information, please refer to :ref:`api_guide_Name` .
diff --git a/test/prim/prim/vjp/eager/test_comp_eager_p_norm_grad.py b/test/prim/prim/vjp/eager/test_comp_eager_p_norm_grad.py
diff --git a/test/prim/prim/vjp/eager/test_comp_eager_reshape_double_grad.py b/test/prim/prim/vjp/eager/test_comp_eager_reshape_double_grad.py
