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

zeroRains · web-flow · commit 555ad25dd236 · 2024-07-15T09:59:53.000+08:00
* surport dynamic shape for  the softmax_grad

* add the configure
diff --git a/paddle/fluid/primitive/rule/vjp/details.h b/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);
     }
   }
 }
diff --git a/python/paddle/autograd/backward_utils.py b/python/paddle/autograd/backward_utils.py
@@ -53,6 +53,7 @@
     "pd_op.divide",
     "pd_op.pow",
     "pd_op.elementwise_pow",
+    "pd_op.softmax",
 ]
 
 
diff --git a/test/prim/pir_prim/test_prim_sub_graph_backward_dynamic_shape.py b/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()

Original file line number	Diff line number	Diff line change
`@@ -1084,7 +1084,8 @@ void softmax_grad(const Tensor& out,`
`1084`	`1084`	`set_output<T>(tmp_x_grad, x_grad);`
`1085`	`1085`	`}`
`1086`	`1086`	`} else {`
`1087`		`- set_output<T>(out_grad * 0.0, x_grad);`
	`1087`	`+ Tensor zeros = full_scalar<T>(0.0, out.dtype());`
	`1088`	`+ set_output<T>(out_grad * zeros, x_grad);`
`1088`	`1089`	`}`
`1089`	`1090`	`}`
`1090`	`1091`	`}`
Original file line number	Diff line number	Diff line change
`@@ -53,6 +53,7 @@`
`53`	`53`	`"pd_op.divide",`
`54`	`54`	`"pd_op.pow",`
`55`	`55`	`"pd_op.elementwise_pow",`
	`56`	`+ "pd_op.softmax",`
`56`	`57`	`]`
`57`	`58`
`58`	`59`