[Prim][PIR] Surport dynamic shape for softmax_grad (#65961)

* surport dynamic shape for  the softmax_grad

* add the configure

Author: zeroRains

paddle/fluid/primitive/rule/vjp/details.h

@@ -1084,7 +1084,8 @@ void softmax_grad(const Tensor& out,
         set_output<T>(tmp_x_grad, x_grad);
       }
     } else {
-      set_output<T>(out_grad * 0.0, x_grad);
+      Tensor zeros = full_scalar<T>(0.0, out.dtype());
+      set_output<T>(out_grad * zeros, x_grad);
     }
   }
 }

python/paddle/autograd/backward_utils.py

@@ -53,6 +53,7 @@
     "pd_op.divide",
     "pd_op.pow",
     "pd_op.elementwise_pow",
+    "pd_op.softmax",
 ]
 
 

test/prim/pir_prim/test_prim_sub_graph_backward_dynamic_shape.py

@@ -115,6 +115,10 @@ def pow_net(x):
     return paddle.pow(x, 3.2)
 
 
+def softmax_net(x):
+    return paddle.nn.functional.softmax(x, axis=-1)
+
+
 def apply_to_static(net, use_cinn, input_spec=None):
     build_strategy = paddle.static.BuildStrategy()
     build_strategy.build_cinn_pass = use_cinn
@@ -1330,5 +1334,41 @@ def setUp(self):
         self.tol = 1e-6
 
 
+class TestPrimSoftmaxWithGrad1(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = softmax_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimSoftmaxWithGrad2(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        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 = softmax_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimSoftmaxWithGrad3(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [30, 200, None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = softmax_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
 if __name__ == "__main__":
     unittest.main()