[Prim] Fix composite subtract double grad in broadcast case (#63697)

HydrogenSulfate · web-flow · commit 20a11588dad8 · 2024-04-22T15:20:33.000+08:00
* Fix broadcast case for subtract_double_grad

* Fix broadcast case for subtract_double_grad

* add UT
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
@@ -773,9 +773,85 @@ void subtract_double_grad(const Tensor& y,
     if (grad_x_grad && grad_y_grad) {
       set_output<T>(grad_x_grad.get() - grad_y_grad.get(), grad_out_grad);
     } else if (grad_x_grad) {
-      by_pass<T>(grad_x_grad.get(), grad_out_grad);
+      if (grad_x_grad.get().dims() != grad_out.dims()) {
+        // broad cast grad_x_grad to grad_out
+        auto grad_x_grad_dims = common::vectorize(grad_x_grad.get().dims());
+        auto grad_out_dims = common::vectorize(grad_out.dims());
+        auto broadcast_dims = grad_x_grad_dims;
+        // reshape to same dims
+        bool need_reshape = false;
+        if (grad_out_dims.size() > grad_x_grad_dims.size()) {
+          need_reshape = true;
+          for (size_t i = 0; i < grad_out_dims.size() - grad_x_grad_dims.size();
+               ++i) {
+            broadcast_dims.insert(broadcast_dims.begin(), 1);
+          }
+        }
+        // tile if needed
+        auto repeat_times = broadcast_dims;
+        bool need_tile = false;
+        for (size_t i = 0; i < broadcast_dims.size(); ++i) {
+          if (grad_out_dims[i] > 1 && broadcast_dims[i] == 1) {
+            repeat_times[i] = grad_out_dims[i];
+            need_tile = true;
+          } else {
+            repeat_times[i] = 1;
+          }
+        }
+        if (need_reshape && need_tile) {
+          set_output<T>(tile<T>(reshape<T>(grad_x_grad.get(), broadcast_dims),
+                                repeat_times),
+                        grad_out_grad);
+        } else if (need_reshape) {
+          set_output<T>(reshape<T>(grad_x_grad.get(), broadcast_dims),
+                        grad_out_grad);
+        } else if (need_tile) {
+          set_output<T>(tile<T>(grad_x_grad.get(), repeat_times),
+                        grad_out_grad);
+        }
+      } else {
+        by_pass<T>(grad_x_grad.get(), grad_out_grad);
+      }
     } else if (grad_y_grad) {
-      by_pass<T>(-grad_y_grad.get(), grad_out_grad);
+      if (grad_y_grad.get().dims() != grad_out.dims()) {
+        // broad cast grad_y_grad to grad_out
+        auto grad_y_grad_dims = common::vectorize(grad_y_grad.get().dims());
+        auto grad_out_dims = common::vectorize(grad_out.dims());
+        auto broadcast_dims = grad_y_grad_dims;
+        // reshape to same dims
+        bool need_reshape = false;
+        if (grad_out_dims.size() > grad_y_grad_dims.size()) {
+          need_reshape = true;
+          for (size_t i = 0; i < grad_out_dims.size() - grad_y_grad_dims.size();
+               ++i) {
+            broadcast_dims.insert(broadcast_dims.begin(), 1);
+          }
+        }
+        // tile if needed
+        auto repeat_times = broadcast_dims;
+        bool need_tile = false;
+        for (size_t i = 0; i < broadcast_dims.size(); ++i) {
+          if (grad_out_dims[i] > 1 && broadcast_dims[i] == 1) {
+            repeat_times[i] = grad_out_dims[i];
+            need_tile = true;
+          } else {
+            repeat_times[i] = 1;
+          }
+        }
+        if (need_reshape && need_tile) {
+          set_output<T>(tile<T>(reshape<T>(grad_y_grad.get(), broadcast_dims),
+                                repeat_times),
+                        grad_out_grad);
+        } else if (need_reshape) {
+          set_output<T>(reshape<T>(grad_y_grad.get(), broadcast_dims),
+                        grad_out_grad);
+        } else if (need_tile) {
+          set_output<T>(tile<T>(grad_y_grad.get(), repeat_times),
+                        grad_out_grad);
+        }
+      } else {
+        by_pass<T>(-grad_y_grad.get(), grad_out_grad);
+      }
     } else {
       set_output<T>(
           full<T>(common::vectorize(grad_out.dims()), 0, grad_out.dtype()),
diff --git a/test/legacy_test/test_elementwise_nn_grad.py b/test/legacy_test/test_elementwise_nn_grad.py
@@ -216,6 +216,211 @@ def test_grad(self):
             self.func(p)
 
 
+class TestElementwiseSubBroadcastDoubleGradCheck2(unittest.TestCase):
+    def subtract_wrapper(self, x):
+        return paddle.subtract(x[0], x[1])
+
+    @test_with_pir_api
+    @prog_scope()
+    def func(self, place):
+        # the shape of input variable should be clearly specified, not include -1.
+        shape1 = [2, 1, 4, 5]
+        shape2 = [2, 3, 1, 1]
+        eps = 0.005
+        dtype = np.float64
+
+        x = paddle.static.data('x', shape1, dtype)
+        y = paddle.static.data('y', shape2, dtype)
+        x.persistable = True
+        y.persistable = True
+        out = paddle.subtract(x, y)
+        x_arr = np.random.uniform(-1, 1, shape1).astype(dtype)
+        y_arr = np.random.uniform(-1, 1, shape2).astype(dtype)
+
+        gradient_checker.double_grad_check(
+            [x, y], out, x_init=[x_arr, y_arr], place=place, eps=eps
+        )
+        gradient_checker.double_grad_check_for_dygraph(
+            self.subtract_wrapper,
+            [x, y],
+            out,
+            x_init=[x_arr, y_arr],
+            place=place,
+        )
+
+    def test_grad(self):
+        paddle.enable_static()
+        places = [base.CPUPlace()]
+        if core.is_compiled_with_cuda():
+            places.append(base.CUDAPlace(0))
+        for p in places:
+            self.func(p)
+
+
+class TestElementwiseSubBroadcastDoubleGradCheck3(unittest.TestCase):
+    def subtract_wrapper(self, x):
+        return paddle.subtract(x[0], x[1])
+
+    @test_with_pir_api
+    @prog_scope()
+    def func(self, place):
+        # the shape of input variable should be clearly specified, not include -1.
+        shape1 = [2, 1, 4, 5]
+        shape2 = [1, 1]
+        eps = 0.005
+        dtype = np.float64
+
+        x = paddle.static.data('x', shape1, dtype)
+        y = paddle.static.data('y', shape2, dtype)
+        x.persistable = True
+        y.persistable = True
+        out = paddle.subtract(x, y)
+        x_arr = np.random.uniform(-1, 1, shape1).astype(dtype)
+        y_arr = np.random.uniform(-1, 1, shape2).astype(dtype)
+
+        gradient_checker.double_grad_check(
+            [x, y], out, x_init=[x_arr, y_arr], place=place, eps=eps
+        )
+        gradient_checker.double_grad_check_for_dygraph(
+            self.subtract_wrapper,
+            [x, y],
+            out,
+            x_init=[x_arr, y_arr],
+            place=place,
+        )
+
+    def test_grad(self):
+        paddle.enable_static()
+        places = [base.CPUPlace()]
+        if core.is_compiled_with_cuda():
+            places.append(base.CUDAPlace(0))
+        for p in places:
+            self.func(p)
+
+
+class TestElementwiseSubBroadcastDoubleGradCheck4(unittest.TestCase):
+    def subtract_wrapper(self, x):
+        return paddle.subtract(x[0], x[1])
+
+    @test_with_pir_api
+    @prog_scope()
+    def func(self, place):
+        # the shape of input variable should be clearly specified, not include -1.
+        shape1 = [2, 1, 4, 5]
+        shape2 = []
+        eps = 0.005
+        dtype = np.float64
+
+        x = paddle.static.data('x', shape1, dtype)
+        y = paddle.static.data('y', shape2, dtype)
+        x.persistable = True
+        y.persistable = True
+        out = paddle.subtract(x, y)
+        x_arr = np.random.uniform(-1, 1, shape1).astype(dtype)
+        y_arr = np.random.uniform(-1, 1, shape2).astype(dtype)
+
+        gradient_checker.double_grad_check(
+            [x, y], out, x_init=[x_arr, y_arr], place=place, eps=eps
+        )
+        gradient_checker.double_grad_check_for_dygraph(
+            self.subtract_wrapper,
+            [x, y],
+            out,
+            x_init=[x_arr, y_arr],
+            place=place,
+        )
+
+    def test_grad(self):
+        paddle.enable_static()
+        places = [base.CPUPlace()]
+        if core.is_compiled_with_cuda():
+            places.append(base.CUDAPlace(0))
+        for p in places:
+            self.func(p)
+
+
+class TestElementwiseSubBroadcastDoubleGradCheck5(unittest.TestCase):
+    def subtract_wrapper(self, x):
+        return paddle.subtract(x[0], x[1])
+
+    @test_with_pir_api
+    @prog_scope()
+    def func(self, place):
+        # the shape of input variable should be clearly specified, not include -1.
+        shape1 = [2, 1, 4, 5]
+        shape2 = [4, 1]
+        eps = 0.005
+        dtype = np.float64
+
+        x = paddle.static.data('x', shape1, dtype)
+        y = paddle.static.data('y', shape2, dtype)
+        x.persistable = True
+        y.persistable = True
+        out = paddle.subtract(x, y)
+        x_arr = np.random.uniform(-1, 1, shape1).astype(dtype)
+        y_arr = np.random.uniform(-1, 1, shape2).astype(dtype)
+
+        gradient_checker.double_grad_check(
+            [x, y], out, x_init=[x_arr, y_arr], place=place, eps=eps
+        )
+        gradient_checker.double_grad_check_for_dygraph(
+            self.subtract_wrapper,
+            [x, y],
+            out,
+            x_init=[x_arr, y_arr],
+            place=place,
+        )
+
+    def test_grad(self):
+        paddle.enable_static()
+        places = [base.CPUPlace()]
+        if core.is_compiled_with_cuda():
+            places.append(base.CUDAPlace(0))
+        for p in places:
+            self.func(p)
+
+
+class TestElementwiseSubBroadcastDoubleGradCheck6(unittest.TestCase):
+    def subtract_wrapper(self, x):
+        return paddle.subtract(x[0], x[1])
+
+    @test_with_pir_api
+    @prog_scope()
+    def func(self, place):
+        # the shape of input variable should be clearly specified, not include -1.
+        shape1 = [4, 1, 3]
+        shape2 = [3, 1]
+        eps = 0.005
+        dtype = np.float64
+
+        x = paddle.static.data('x', shape1, dtype)
+        y = paddle.static.data('y', shape2, dtype)
+        x.persistable = True
+        y.persistable = True
+        out = paddle.subtract(x, y)
+        x_arr = np.random.uniform(-1, 1, shape1).astype(dtype)
+        y_arr = np.random.uniform(-1, 1, shape2).astype(dtype)
+
+        gradient_checker.double_grad_check(
+            [x, y], out, x_init=[x_arr, y_arr], place=place, eps=eps
+        )
+        gradient_checker.double_grad_check_for_dygraph(
+            self.subtract_wrapper,
+            [x, y],
+            out,
+            x_init=[x_arr, y_arr],
+            place=place,
+        )
+
+    def test_grad(self):
+        paddle.enable_static()
+        places = [base.CPUPlace()]
+        if core.is_compiled_with_cuda():
+            places.append(base.CUDAPlace(0))
+        for p in places:
+            self.func(p)
+
+
 class TestElementwiseDivDoubleGradCheck(unittest.TestCase):
     def divide_wrapper(self, x):
         return paddle.divide(x[0], x[1])
