[Prim][PIR] Add the dynamic shape test case for scale_grad, square_grad, transpose_grad and swiglu_grad

python/paddle/autograd/backward_utils.py

@@ -68,6 +68,10 @@
     "pd_op.max",
     "pd_op.stack",
     "pd_op.expand",
+    "pd_op.scale",
+    "pd_op.square",
+    "pd_op.transpose",
+    "pd_op.swiglu",
 ]
 
 

test/prim/pir_prim/test_prim_sub_graph_backward_dynamic_shape.py

@@ -217,6 +217,26 @@ def stack_net3(x):
     return paddle.stack(x, axis=0)
 
 
+def scale_net(x):
+    return paddle.scale(x, scale=-2.3)
+
+
+def square_net(x):
+    return paddle.square(x)
+
+
+def transpose_net(x):
+    return paddle.transpose(x, perm=[0, 3, 1, 2])
+
+
+def swiglu_net1(x, y):
+    return paddle.incubate.nn.functional.swiglu(x, y)
+
+
+def swiglu_net2(x):
+    return paddle.incubate.nn.functional.swiglu(x)
+
+
 def apply_to_static(net, use_cinn, input_spec=None):
     build_strategy = paddle.static.BuildStrategy()
     build_strategy.build_cinn_pass = use_cinn
@@ -2214,5 +2234,69 @@ def setUp(self):
         self.tol = 1e-6
 
 
+class TestPrimScaleWithGrad(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [20, 30, 70]
+        self.init_x_shape = [None, None, 70]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = scale_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimSquareWithGrad(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [20, 30, 70]
+        self.init_x_shape = [None, None, 70]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = square_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimTransposeWithGrad(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [20, 30, 50, 70]
+        self.init_x_shape = [None, None, None, 70]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = transpose_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimSwigluWithGrad1(TestPrimBaseOneGradTwoInputs):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.y_shape = [30, 200, 40]
+        self.init_y_shape = [None, None, 40]
+        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.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.net = swiglu_net1
+        self.enable_cinn = False
+        self.tol = 1e-5
+        self.y_without_grad = True
+
+
+class TestPrimSwigluWithGrad2(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [20, 30, 50, 70]
+        self.init_x_shape = [None, None, None, 70]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = swiglu_net2
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
 if __name__ == "__main__":
     unittest.main()