add backward decomp for logsumexp and trunc (#68485)

eggman-1024 · web-flow · commit b4baafae3682 · 2024-10-14T10:18:46.000+08:00
* add backward decomp for logsumexp and trunc

* modified logsumexp_grad

* modified some files

* add prim_op_type and public_python_api for TestLogsumexpBF16Op()

* support dynamic shape for logsumexp_grad

* support dynamic shape for trunc_grad &amp; modify TestLogsumexp_FP16
diff --git a/paddle/fluid/primitive/codegen/gen.py b/paddle/fluid/primitive/codegen/gen.py
@@ -76,6 +76,7 @@
     'sin_grad',
     'cos_grad',
     'tanh_grad',
+    'trunc_grad',
     'square_grad',
 ]
 
@@ -126,6 +127,7 @@
     'unsqueeze_grad',
     'where_grad',
     'logcumsumexp_grad',
+    'logsumexp_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
@@ -2884,6 +2884,106 @@ void logcumsumexp_grad(const Tensor& x,
   }
 }
 
+template <typename T>
+void logsumexp_grad(const Tensor& x,
+                    const Tensor& out,
+                    const Tensor& out_grad,
+                    const IntArray& axis,
+                    bool keepdim,
+                    bool reduce_all,
+                    Tensor* x_grad) {
+  if (x_grad) {
+    int64_t axis_size = axis.size();
+    int64_t x_dim_size = x.dims().size();
+    reduce_all = false;
+
+    if (reduce_all || axis_size == 0 || axis_size == x_dim_size) {
+      reduce_all = true;
+    } else {
+      reduce_all = false;
+    }
+
+    auto x_grad_tmp = Tensor();
+
+    if (has_dynamic_shape(x.shape())) {
+      Tensor x_shape = shape<T>(x);
+      if (x_dim_size == 1) {
+        x_grad_tmp = backend::expand<T>(out_grad, x_shape) * exp<T>(x - out);
+      } else {
+        if (!keepdim) {
+          auto axis_ = std::vector<int64_t>();
+          if (reduce_all) {
+            for (int64_t i = 0; i < x_dim_size; i++) {
+              axis_.push_back(i);
+            }
+          } else {
+            axis_ = axis.GetData();
+            for (int64_t i = 0; i < axis_size; i++) {
+              if (axis[i] < 0) {
+                axis_[i] = axis[i] + x_dim_size;
+              }
+            }
+          }
+
+          auto result_shape = get_unsqueeze_dims<T>(shape<T>(out_grad), axis_);
+          auto out_ = backend::reshape<T>(out, result_shape);
+          auto softmax = exp<T>(x - backend::expand<T>(out_, x_shape));
+
+          auto out_grad_ = backend::reshape<T>(out_grad, result_shape);
+          x_grad_tmp = backend::expand<T>(out_grad_, x_shape) * softmax;
+        } else {
+          x_grad_tmp = backend::expand<T>(out_grad, x_shape) * exp<T>(x - out);
+        }
+      }
+    } else {
+      std::vector<int64_t> x_dim = common::vectorize<int64_t>(x.dims());
+      if (x_dim_size == 1) {
+        x_grad_tmp = expand<T>(out_grad, IntArray(x_dim)) * exp<T>(x - out);
+      } else {
+        if (!keepdim) {
+          auto axis_ = std::vector<int64_t>();
+          if (reduce_all) {
+            for (int64_t i = 0; i < x_dim_size; i++) {
+              axis_.push_back(i);
+            }
+          } else {
+            axis_ = axis.GetData();
+            for (int64_t i = 0; i < axis_size; i++) {
+              if (axis[i] < 0) {
+                axis_[i] = axis[i] + x_dim_size;
+              }
+            }
+          }
+          auto out_shape = get_unsqueeze_dims(out, axis_);
+          auto out_ = reshape<T>(out, out_shape);
+          auto softmax = exp<T>(x - expand<T>(out_, IntArray(x_dim)));
+
+          auto out_grad_shape = get_unsqueeze_dims(out_grad, axis_);
+          auto out_grad_ = reshape<T>(out_grad, out_grad_shape);
+          x_grad_tmp = expand<T>(out_grad_, IntArray(x_dim)) * softmax;
+        } else {
+          x_grad_tmp = expand<T>(out_grad, IntArray(x_dim)) * exp<T>(x - out);
+        }
+      }
+    }
+    set_output<T>(x_grad_tmp, x_grad);
+  }
+}
+
+template <typename T>
+void trunc_grad(const Tensor& out_grad, Tensor* x_grad) {
+  Tensor zero;
+  if (x_grad) {
+    if (has_dynamic_shape(out_grad.shape())) {
+      zero = backend::full_with_tensor<T>(
+          shape<T>(out_grad), 0.0, out_grad.dtype());
+    } else {
+      zero = full<T>(out_grad.shape(), 0.0, out_grad.dtype());
+    }
+    set_output<T>(zero, x_grad);
+  }
+}
+
 }  // namespace details
 }  // namespace primitive
 }  // namespace paddle
diff --git a/python/paddle/autograd/backward_utils.py b/python/paddle/autograd/backward_utils.py
@@ -55,6 +55,7 @@
     "pd_op.leaky_relu",
     "pd_op.log",
     "pd_op.logcumsumexp",
+    "pd_op.logsumexp",
     "pd_op.matmul",
     "pd_op.max",
     "pd_op.maximum",
@@ -88,8 +89,9 @@
     "pd_op.swish",
     "pd_op.tanh",
     "pd_op.topk",
-    "pd_op.unsqueeze",
     "pd_op.transpose",
+    "pd_op.trunc",
+    "pd_op.unsqueeze",
     "pd_op.where",
 ]
 
diff --git a/test/legacy_test/test_logsumexp.py b/test/legacy_test/test_logsumexp.py
@@ -57,7 +57,9 @@ def logsumexp_ref_grad(x):
 class TestLogsumexp(OpTest):
     def setUp(self):
         self.op_type = 'logsumexp'
+        self.prim_op_type = "prim"
         self.python_api = logsumexp_wrapper
+        self.public_python_api = logsumexp_wrapper
         self.shape = [2, 3, 4, 5]
         self.dtype = 'float64'
         self.axis = [-1]
@@ -87,7 +89,10 @@ def set_attrs_addition(self):
         pass
 
     def test_check_output(self):
-        self.check_output(check_pir=True)
+        self.check_output(
+            check_pir=True,
+            check_prim_pir=True,
+        )
 
     def test_check_grad(self):
         self.check_grad(
@@ -96,6 +101,7 @@ def test_check_grad(self):
             user_defined_grads=self.user_defined_grads,
             user_defined_grad_outputs=self.user_defined_grad_outputs,
             check_pir=True,
+            check_prim_pir=True,
         )
 
     def calc_grad(self):
@@ -165,24 +171,25 @@ def set_attrs(self):
         self.dtype = 'float16'
 
     def test_check_output(self):
-        ref_x = self.inputs['X'].astype(np.float32)
-        out_ref = ref_logsumexp(ref_x)
-        paddle.disable_static()
-        x = self.inputs['X'].astype(np.float16)
-        tensor_x = paddle.to_tensor(x)
-        out_pad = logsumexp_wrapper(tensor_x)
-        paddle.enable_static()
-        np.testing.assert_allclose(
-            out_pad.numpy(), out_ref, rtol=1e-03, atol=1e-08
+        place = core.CUDAPlace(0)
+        self.check_output_with_place(
+            place,
+            check_pir=True,
+            check_prim_pir=True,
         )
 
     def test_check_grad(self):
-        self.__class__.dtype = self.dtype
-        ref_x = self.inputs['X'].astype(np.float32)
-        ref_x_grad = logsumexp_ref_grad(ref_x)
-        x = self.inputs['X'].astype(np.float16)
-        x_grad = logsumexp_op_grad(x)
-        np.testing.assert_allclose(x_grad, ref_x_grad, rtol=1e-03, atol=1e-05)
+        place = core.CUDAPlace(0)
+        self.check_grad_with_place(
+            place,
+            ['X'],
+            'Out',
+            check_pir=True,
+            check_prim_pir=True,
+        )
+
+    def set_attrs_addition(self):
+        pass
 
 
 @unittest.skipIf(
@@ -193,7 +200,9 @@ def test_check_grad(self):
 class TestLogsumexpBF16Op(TestLogsumexp):
     def setUp(self):
         self.op_type = 'logsumexp'
+        self.prim_op_type = "prim"
         self.python_api = logsumexp_wrapper
+        self.public_python_api = logsumexp_wrapper
         self.dtype = np.uint16
         self.shape = [2, 3, 4, 5]
         self.axis = [-1]
@@ -213,11 +222,21 @@ def setUp(self):
 
     def test_check_output(self):
         place = core.CUDAPlace(0)
-        self.check_output_with_place(place, check_pir=True)
+        self.check_output_with_place(
+            place,
+            check_pir=True,
+            check_prim_pir=True,
+        )
 
     def test_check_grad(self):
         place = core.CUDAPlace(0)
-        self.check_grad_with_place(place, ['X'], 'Out', check_pir=True)
+        self.check_grad_with_place(
+            place,
+            ['X'],
+            'Out',
+            check_pir=True,
+            check_prim_pir=True,
+        )
 
     def set_attrs(self):
         pass
@@ -227,7 +246,6 @@ def set_attrs_addition(self):
 
 
 class TestLogsumexpError(unittest.TestCase):
-
     def test_errors(self):
         with paddle.static.program_guard(paddle.static.Program()):
             self.assertRaises(TypeError, paddle.logsumexp, 1)
diff --git a/test/legacy_test/test_trunc_op.py b/test/legacy_test/test_trunc_op.py
@@ -26,7 +26,9 @@
 class TestTruncOp(OpTest):
     def setUp(self):
         self.op_type = "trunc"
+        self.prim_op_type = "prim"
         self.python_api = paddle.trunc
+        self.public_python_api = paddle.trunc
         self.init_dtype_type()
         np.random.seed(2021)
         self.inputs = {'X': np.random.random((20, 20)).astype(self.dtype)}
@@ -36,10 +38,19 @@ def init_dtype_type(self):
         self.dtype = np.float64
 
     def test_check_output(self):
-        self.check_output(check_pir=True)
+        self.check_output(
+            check_pir=True,
+            check_prim_pir=True,
+        )
 
     def test_check_grad(self):
-        self.check_grad(['X'], 'Out', numeric_grad_delta=1e-5, check_pir=True)
+        self.check_grad(
+            ['X'],
+            'Out',
+            numeric_grad_delta=1e-5,
+            check_pir=True,
+            check_prim_pir=True,
+        )
 
 
 class TestFloatTruncOp(TestTruncOp):
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
@@ -39,6 +39,22 @@ def logcumsumexp_net3(x):
     return paddle.logcumsumexp(x, axis=-1)
 
 
+def logsumexp_net1(x):
+    return paddle.logsumexp(x)
+
+
+def logsumexp_net2(x):
+    return paddle.logsumexp(x, keepdim=False)
+
+
+def logsumexp_net3(x):
+    return paddle.logsumexp(x, axis=-1, keepdim=False)
+
+
+def logsumexp_net4(x):
+    return paddle.logsumexp(x, axis=[0, 2], keepdim=False)
+
+
 def matmul_net(x, y):
     return paddle.matmul(x, y)
 
@@ -143,6 +159,71 @@ def setUp(self):
         self.tol = 1e-6
 
 
+class TestPrimLogsumexpWithGrad1(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.logsumexp_grad"
+        self.dtype = "float32"
+        self.x_shape = [1000]
+        self.init_x_shape = [None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = logsumexp_net1
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimLogsumexpWithGrad2(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.logsumexp_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = logsumexp_net2
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimLogsumexpWithGrad3(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.logsumexp_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = logsumexp_net1
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimLogsumexpWithGrad4(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.logsumexp_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = logsumexp_net3
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimLogsumexpWithGrad5(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.logsumexp_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = logsumexp_net4
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
 class TestPrimMatmulWithGrad1(TestPrimTwoWithGrad):
     def setUp(self):
         np.random.seed(2023)
diff --git a/test/prim/pir_prim/test_prim_sub_graph_pqrst_backward_dynamic_shape.py b/test/prim/pir_prim/test_prim_sub_graph_pqrst_backward_dynamic_shape.py

Original file line number	Diff line number	Diff line change
`@@ -76,6 +76,7 @@`
`76`	`76`	`'sin_grad',`
`77`	`77`	`'cos_grad',`
`78`	`78`	`'tanh_grad',`
	`79`	`+ 'trunc_grad',`
`79`	`80`	`'square_grad',`
`80`	`81`	`]`
`81`	`82`
`@@ -126,6 +127,7 @@`
`126`	`127`	`'unsqueeze_grad',`
`127`	`128`	`'where_grad',`
`128`	`129`	`'logcumsumexp_grad',`
	`130`	`+ 'logsumexp_grad',`
`129`	`131`	`]`
`130`	`132`
`131`	`133`	`# whole vjp list of primitive op vjp`