[Prim][PIR] Add the dynamic shape test case for unsqueeze_grad ans squeeze_grad (#68076)

zeroRains · web-flow · commit 51d82f191ed1 · 2024-09-18T09:26:34.000+08:00
* support dynamic shape for unsqueeze_grad

* add the dynamic shape  test case for squeeze_grad

* fix the op name
diff --git a/python/paddle/autograd/backward_utils.py b/python/paddle/autograd/backward_utils.py
@@ -75,11 +75,13 @@
     "pd_op.split",
     "pd_op.sqrt",
     "pd_op.square",
+    "pd_op.squeeze",
     "pd_op.stack",
     "pd_op.subtract",
     "pd_op.sum",
     "pd_op.swiglu",
     "pd_op.tanh",
+    "pd_op.unsqueeze",
     "pd_op.transpose",
     "pd_op.where",
 ]
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
@@ -276,7 +276,7 @@ def setUp(self):
 class TestPrimBatchNormWithGrad1(TestPrimThreeWithGrad):
     def setUp(self):
         np.random.seed(2023)
-        self.op_name = "pd_op.batch_norm_grad_"
+        self.op_name = "pd_op.batch_norm_grad"
         self.dtype = "float32"
         self.x_shape = [30, 40, 50, 60]
         self.init_x_shape = [None, 40, None, None]
@@ -295,7 +295,7 @@ def setUp(self):
 class TestPrimBatchNormWithGrad2(TestPrimThreeWithGrad):
     def setUp(self):
         np.random.seed(2023)
-        self.op_name = "pd_op.batch_norm_grad_"
+        self.op_name = "pd_op.batch_norm_grad"
         self.dtype = "float32"
         self.x_shape = [30, 40, 50, 60]
         self.init_x_shape = [None, 40, None, None]
@@ -314,7 +314,7 @@ def setUp(self):
 class TestPrimBatchNormWithGrad3(TestPrimThreeWithGrad):
     def setUp(self):
         np.random.seed(2023)
-        self.op_name = "pd_op.batch_norm_grad_"
+        self.op_name = "pd_op.batch_norm_grad"
         self.dtype = "float32"
         self.x_shape = [30, 40, 50, 60]
         self.init_x_shape = [None, None, None, 60]
@@ -333,7 +333,7 @@ def setUp(self):
 class TestPrimBatchNormWithGrad4(TestPrimThreeWithGrad):
     def setUp(self):
         np.random.seed(2023)
-        self.op_name = "pd_op.batch_norm_grad_"
+        self.op_name = "pd_op.batch_norm_grad"
         self.dtype = "float32"
         self.x_shape = [30, 40, 50, 60]
         self.init_x_shape = [None, 40, None, None]
@@ -352,7 +352,7 @@ def setUp(self):
 class TestPrimBatchNormWithGrad5(TestPrimThreeWithGrad):
     def setUp(self):
         np.random.seed(2023)
-        self.op_name = "pd_op.batch_norm_grad_"
+        self.op_name = "pd_op.batch_norm_grad"
         self.dtype = "float32"
         self.x_shape = [30, 40, 50, 60]
         self.init_x_shape = [None, None, None, 60]
@@ -371,7 +371,7 @@ def setUp(self):
 class TestPrimBatchNormWithGrad6(TestPrimThreeWithGrad):
     def setUp(self):
         np.random.seed(2023)
-        self.op_name = "pd_op.batch_norm_grad_"
+        self.op_name = "pd_op.batch_norm_grad"
         self.dtype = "float32"
         self.x_shape = [30, 40, 50, 60]
         self.init_x_shape = [None, 40, None, None]
@@ -390,7 +390,7 @@ def setUp(self):
 class TestPrimBatchNormWithGrad7(TestPrimThreeWithGrad):
     def setUp(self):
         np.random.seed(2023)
-        self.op_name = "pd_op.batch_norm_grad_"
+        self.op_name = "pd_op.batch_norm_grad"
         self.dtype = "float32"
         self.x_shape = [30, 40, 50, 60]
         self.init_x_shape = [None, None, None, 60]
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
@@ -108,6 +108,10 @@ def square_net(x):
     return paddle.square(x)
 
 
+def squeeze_net(x):
+    return paddle.squeeze(x, axis=[0, -2])
+
+
 def stack_net1(x):
     y = x + 1
     return paddle.stack([x, y], axis=-1)
@@ -545,6 +549,45 @@ def setUp(self):
         self.tol = 1e-6
 
 
+class TestPrimSqueezeWithGrad1(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.op_name = "pd_op.squeeze_grad"
+        self.dtype = "float32"
+        self.x_shape = [1, 20, 1, 30]
+        self.init_x_shape = [None, None, None, 30]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = squeeze_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimSqueezeWithGrad2(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.op_name = "pd_op.squeeze_grad"
+        self.dtype = "float32"
+        self.x_shape = [1, 20, 1, 30]
+        self.init_x_shape = [None, 20, None, 30]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = squeeze_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimSqueezeWithGrad3(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.op_name = "pd_op.squeeze_grad"
+        self.dtype = "float32"
+        self.x_shape = [1, 20, 1, 30]
+        self.init_x_shape = [1, None, 1, None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = squeeze_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
 class TestPrimStackWithGrad1(TestPrimBaseWithGrad):
     def setUp(self):
         np.random.seed(2023)
diff --git a/test/prim/pir_prim/test_prim_sub_graph_uvwxyz_backward_dynamic_shape.py b/test/prim/pir_prim/test_prim_sub_graph_uvwxyz_backward_dynamic_shape.py
@@ -16,16 +16,34 @@
 
 import numpy as np
 from test_prim_sub_graph_backward_dynamic_shape import (
+    TestPrimBaseWithGrad,
     TestPrimTwoWithGrad,
 )
 
 import paddle
 
 
+def unsqueeze_net(x):
+    return paddle.unsqueeze(x, axis=[1, 2])
+
+
 def where_net(x, y):
     return paddle.where(x > y, x, y)
 
 
+class TestPrimUnsqueezeWithGrad(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.unsqueeze_grad"
+        self.dtype = "float32"
+        self.x_shape = [20, 30, 40]
+        self.init_x_shape = [None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = unsqueeze_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
 class TestPrimWhereWithGrad1(TestPrimTwoWithGrad):
     def setUp(self):
         np.random.seed(2023)
