decomp logcumsumexp_grad (#68346)

cubehan3 · web-flow · commit f874c6750e5c · 2024-09-24T09:57:58.000+08:00
diff --git a/paddle/fluid/primitive/codegen/gen.py b/paddle/fluid/primitive/codegen/gen.py
@@ -122,6 +122,7 @@
     'topk_grad',
     'unsqueeze_grad',
     'where_grad',
+    'logcumsumexp_grad',
 ]
 
 # whole vjp list of primitive op vjp
diff --git a/paddle/fluid/primitive/rule/vjp/details.h b/paddle/fluid/primitive/rule/vjp/details.h
@@ -2808,6 +2808,104 @@ void atan_grad(const Tensor& x, const Tensor& out_grad, Tensor* x_grad) {
   }
 }
 
+template <typename T>
+void logcumsumexp_grad(const Tensor& x,
+                       const Tensor& out,
+                       const Tensor& out_grad,
+                       int axis,
+                       bool flatten,
+                       bool exclusive,
+                       bool reverse,
+                       Tensor* x_grad) {
+  if (x_grad) {
+    reverse = !reverse;
+    Tensor tmp, lowest, x_grad_tmp;
+    Tensor x_cast = x;
+    Tensor out_cast = out;
+    Tensor out_grad_cast = out_grad;
+    bool need_cast = is_half_dtype(x.dtype());
+
+    if (need_cast) {
+      x_cast = cast<T>(x, DataType::FLOAT32);
+      out_cast = cast<T>(out, DataType::FLOAT32);
+      out_grad_cast = cast<T>(out_grad, DataType::FLOAT32);
+    }
+    const Tensor out_grad_log = log<T>(abs<T>(out_grad_cast));
+    auto out_grad_dtype = out_grad_cast.dtype();
+
+    if (has_dynamic_shape(x_cast.shape()) ||
+        has_dynamic_shape(out_grad_cast.shape())) {
+      const Tensor x_shape = shape<T>(x_cast);
+      const Tensor out_grad_shape = shape<T>(out_grad_cast);
+      const Tensor reshape_x = backend::reshape<T>(x_cast, out_grad_shape);
+
+      if (out_grad_dtype == DataType::FLOAT32) {
+        lowest =
+            backend::full_with_tensor<T>(out_grad_shape,
+                                         std::numeric_limits<float>::lowest(),
+                                         out_grad_dtype);
+      } else if (out_grad_dtype == DataType::FLOAT64) {
+        lowest =
+            backend::full_with_tensor<T>(out_grad_shape,
+                                         std::numeric_limits<double>::lowest(),
+                                         out_grad_dtype);
+      }
+      const Tensor zero =
+          backend::full_with_tensor<T>(out_grad_shape, 0.0, out_grad_dtype);
+
+      // compute positive
+      Tensor out_grad_pos =
+          where<T>(out_grad_cast > zero, out_grad_log, lowest);
+      tmp = out_grad_pos - out_cast;
+      out_grad_pos = logcumsumexp<T>(tmp, axis, flatten, exclusive, reverse);
+      out_grad_pos = exp<T>(out_grad_pos + reshape_x);
+
+      // compute negative
+      Tensor out_grad_neg =
+          where<T>(out_grad_cast < zero, out_grad_log, lowest);
+      tmp = out_grad_neg - out_cast;
+      out_grad_neg = logcumsumexp<T>(tmp, axis, flatten, exclusive, reverse);
+      out_grad_neg = exp<T>(out_grad_neg + reshape_x);
+
+      x_grad_tmp = backend::reshape<T>(out_grad_pos - out_grad_neg, x_shape);
+    } else {
+      const Tensor reshape_x = reshape<T>(x_cast, out_grad_cast.shape());
+      if (out_grad_dtype == DataType::FLOAT32) {
+        lowest = full<T>(out_grad_cast.shape(),
+                         std::numeric_limits<float>::lowest(),
+                         out_grad_dtype);
+      } else if (out_grad_dtype == DataType::FLOAT64) {
+        lowest = full<T>(out_grad_cast.shape(),
+                         std::numeric_limits<double>::lowest(),
+                         out_grad_dtype);
+      }
+      const Tensor zero = full<T>(out_grad_cast.shape(), 0.0, out_grad_dtype);
+
+      // compute positive
+      Tensor out_grad_pos =
+          where<T>(out_grad_cast > zero, out_grad_log, lowest);
+      tmp = out_grad_pos - out_cast;
+      out_grad_pos = logcumsumexp<T>(tmp, axis, flatten, exclusive, reverse);
+      out_grad_pos = exp<T>(out_grad_pos + reshape_x);
+
+      // compute negative
+      Tensor out_grad_neg =
+          where<T>(out_grad_cast < zero, out_grad_log, lowest);
+      tmp = out_grad_neg - out_cast;
+      out_grad_neg = logcumsumexp<T>(tmp, axis, flatten, exclusive, reverse);
+      out_grad_neg = exp<T>(out_grad_neg + reshape_x);
+
+      x_grad_tmp = reshape<T>(out_grad_pos - out_grad_neg, x_cast.shape());
+    }
+
+    if (need_cast) {
+      set_output<T>(cast<T>(x_grad_tmp, x.dtype()), x_grad);
+    } else {
+      set_output<T>(x_grad_tmp, x_grad);
+    }
+  }
+}
+
 }  // namespace details
 }  // namespace primitive
 }  // namespace paddle
diff --git a/python/paddle/autograd/backward_utils.py b/python/paddle/autograd/backward_utils.py
@@ -51,6 +51,7 @@
     "pd_op.hardswish",
     "pd_op.leaky_relu",
     "pd_op.log",
+    "pd_op.logcumsumexp",
     "pd_op.matmul",
     "pd_op.max",
     "pd_op.maximum",
diff --git a/test/legacy_test/test_logcumsumexp_op.py b/test/legacy_test/test_logcumsumexp_op.py
@@ -234,7 +234,9 @@ class BaseTestCases:
     class BaseOpTest(OpTest):
         def setUp(self):
             self.op_type = "logcumsumexp"
+            self.prim_op_type = "prim"
             self.python_api = logcumsumexp_wrapper
+            self.public_python_api = logcumsumexp_wrapper
             input, attrs = self.input_and_attrs()
             self.inputs = {'X': input}
             self.attrs = attrs
@@ -257,6 +259,7 @@ def test_check_grad(self):
                     )
                 ],
                 check_pir=True,
+                check_prim_pir=True,
             )
 
         def input_and_attrs(self):
@@ -334,8 +337,10 @@ def test_main(self):
 class TestLogcumsumexpBF16Op(OpTest):
     def setUp(self):
         self.op_type = 'logcumsumexp'
+        self.prim_op_type = 'prim'
         self.dtype = np.uint16
         self.python_api = logcumsumexp_wrapper
+        self.public_python_api = logcumsumexp_wrapper
         x = np.arange(100, dtype=np.float64).reshape(10, 10)
         output = np_logcumsumexp(x)
         self.inputs = {'X': convert_float_to_uint16(x)}
@@ -371,6 +376,7 @@ def test_check_grad(self):
             numeric_grad_delta=0.5,
             max_relative_error=0.5,
             check_pir=True,
+            check_prim_pir=True,
         )
 
 
diff --git a/test/prim/pir_prim/test_prim_sub_graph_klmno_backward_dynamic_shape.py b/test/prim/pir_prim/test_prim_sub_graph_klmno_backward_dynamic_shape.py
@@ -27,6 +27,18 @@ def leaky_relu_net(x):
     return paddle.nn.functional.leaky_relu(x)
 
 
+def logcumsumexp_net1(x):
+    return paddle.logcumsumexp(x)
+
+
+def logcumsumexp_net2(x):
+    return paddle.logcumsumexp(x, axis=0)
+
+
+def logcumsumexp_net3(x):
+    return paddle.logcumsumexp(x, axis=-1)
+
+
 def matmul_net(x, y):
     return paddle.matmul(x, y)
 
@@ -92,6 +104,45 @@ def setUp(self):
         self.tol = 1e-6
 
 
+class TestPrimLogcumsumexpWithGrad1(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.logcumsumexp_grad"
+        self.dtype = "float32"
+        self.x_shape = [10, 10, 10]
+        self.init_x_shape = [None, None, None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = logcumsumexp_net1
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimLogcumsumexpWithGrad2(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.logcumsumexp_grad"
+        self.dtype = "float32"
+        self.x_shape = [10, 10, 10]
+        self.init_x_shape = [None, None, None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = logcumsumexp_net2
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimLogcumsumexpWithGrad3(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.logcumsumexp_grad"
+        self.dtype = "float32"
+        self.x_shape = [10, 10, 10]
+        self.init_x_shape = [None, None, None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = logcumsumexp_net3
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
 class TestPrimMatmulWithGrad1(TestPrimTwoWithGrad):
     def setUp(self):
         np.random.seed(2023)

Original file line number	Diff line number	Diff line change
`@@ -122,6 +122,7 @@`
`122`	`122`	`'topk_grad',`
`123`	`123`	`'unsqueeze_grad',`
`124`	`124`	`'where_grad',`
	`125`	`+ 'logcumsumexp_grad',`
`125`	`126`	`]`
`126`	`127`
`127`	`128`	`# whole vjp list of primitive op vjp`