[Prim][Pir] Add backward decomp and support dynamic for angle_grad(not complex) (#70080)

eggman-1024 · web-flow · commit a35a80409542 · 2024-12-27T19:43:09.000+08:00
* add backward decomp and support dynamic for angle_grad(except complex type)

* fix bugs

* rm complex type judgment and add place()

* .

* add 'angle_grad' in decomp_interface_gen_op_list.py
diff --git a/paddle/fluid/pir/dialect/op_generator/decomp_interface_gen_op_list.py b/paddle/fluid/pir/dialect/op_generator/decomp_interface_gen_op_list.py
@@ -93,6 +93,7 @@
 GENERATE_IMPL_VJP = [
     'abs_grad',
     'add_grad',
+    'angle_grad',
     'bce_loss_grad',
     'cos_grad',
     'divide_grad',
diff --git a/paddle/fluid/primitive/codegen/decomp_vjp_gen.py b/paddle/fluid/primitive/codegen/decomp_vjp_gen.py
@@ -72,6 +72,7 @@
     'add_grad',
     'amax_grad',
     'amin_grad',
+    'angle_grad',
     'argsort_grad',
     'assign_grad',
     'atan_grad',
diff --git a/paddle/fluid/primitive/decomp_rule/decomp_vjp/details.h b/paddle/fluid/primitive/decomp_rule/decomp_vjp/details.h
@@ -3630,6 +3630,23 @@ void p_norm_grad(const Tensor& x,
   }
 }
 
+template <typename T>
+void angle_grad(const Tensor& x, const Tensor& out_grad, Tensor* x_grad) {
+  if (x_grad) {
+    Tensor cast_x = ConvertToMT<T>(x);
+    Tensor zero_tensor;
+    if (has_dynamic_shape(cast_x.shape())) {
+      const Tensor x_shape = shape64<T>(cast_x);
+      zero_tensor = backend::full_with_tensor<T>(
+          x_shape, 0.0, cast_x.dtype(), cast_x.place());
+    } else {
+      zero_tensor = full<T>(cast_x.shape(), 0, cast_x.dtype(), cast_x.place());
+    }
+
+    set_output<T>(ConvertToOrig<T>(zero_tensor, x.dtype()), x_grad);
+  }
+}
+
 }  // namespace details
 }  // namespace primitive
 }  // namespace paddle
diff --git a/python/paddle/autograd/backward_utils.py b/python/paddle/autograd/backward_utils.py
@@ -33,6 +33,7 @@
     "pd_op.add",
     "pd_op.amax",
     "pd_op.amin",
+    "pd_op.angle",
     "pd_op.argsort",
     "pd_op.assign",
     "pd_op.batch_norm_",
diff --git a/test/legacy_test/test_angle_op.py b/test/legacy_test/test_angle_op.py
@@ -42,6 +42,8 @@ class TestAngleOpFloat(OpTest):
     def setUp(self):
         self.op_type = "angle"
         self.python_api = paddle.angle
+        self.prim_op_type = "prim"
+        self.public_python_api = paddle.angle
         self.dtype = "float64"
         self.x = np.linspace(-5, 5, 101).astype(self.dtype)
         out_ref = np.angle(self.x)
@@ -59,13 +61,16 @@ def test_check_grad(self):
                 angle_grad(self.x, np.ones_like(self.x) / self.x.size)
             ],
             check_pir=True,
+            check_prim_pir=True,
         )
 
 
 class TestAngleFP16Op(TestAngleOpFloat):
     def setUp(self):
         self.op_type = "angle"
         self.python_api = paddle.angle
+        self.prim_op_type = "prim"
+        self.public_python_api = paddle.angle
         self.dtype = "float16"
         self.x = np.linspace(-5, 5, 101).astype(self.dtype)
         out_ref = np.angle(self.x)
@@ -82,6 +87,8 @@ class TestAngleBF16Op(OpTest):
     def setUp(self):
         self.op_type = "angle"
         self.python_api = paddle.angle
+        self.prim_op_type = "prim"
+        self.public_python_api = paddle.angle
         self.dtype = np.uint16
         self.np_dtype = np.float32
         self.x = np.linspace(-5, 5, 101).astype(self.np_dtype)
@@ -107,6 +114,7 @@ def test_check_grad(self):
                 angle_grad(self.x, np.ones_like(self.x) / self.x.size)
             ],
             check_pir=True,
+            check_prim_pir=True,
         )
 
 
diff --git a/test/prim/pir_prim/test_prim_sub_graph_abcde_backward_dynamic_shape.py b/test/prim/pir_prim/test_prim_sub_graph_abcde_backward_dynamic_shape.py
@@ -71,6 +71,10 @@ def amin_net5(x):
     return paddle.amin(x, axis=[-1, 0], keepdim=False)
 
 
+def angle_net(x):
+    return paddle.angle(x)
+
+
 def argsort_net1(x):
     return paddle.argsort(x, axis=-1)
 
@@ -495,6 +499,19 @@ def setUp(self):
         self.tol = 1e-6
 
 
+class TestPrimAngleWithGrad(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.angle_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 = angle_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
 class TestPrimArgsortWithGrad1(TestPrimBaseWithGrad):
     def setUp(self):
         np.random.seed(2024)
