[Prim][PIR] Support dynamic shape for relu_grad

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

@@ -1287,12 +1287,9 @@ void masked_select_grad(const Tensor& x,
 template <typename T>
 void relu_grad(const Tensor& out, const Tensor& out_grad, Tensor* x_grad) {
   if (x_grad) {
-    auto condition = greater_than<T>(
-        out, full<T>(common::vectorize(out.dims()), 0.0, out.dtype()));
-    auto res =
-        where<T>(condition,
-                 out_grad,
-                 full<T>(common::vectorize(out.dims()), 0.0, out.dtype()));
+    Tensor zeros = full_scalar<T>(0.0, out.dtype());
+    auto mask = greater_than<T>(out, zeros);
+    auto res = cast<T>(mask, out.dtype()) * out_grad;
     set_output<T>(res, x_grad);
   }
 }

python/paddle/autograd/backward_utils.py

@@ -48,6 +48,8 @@
     "pd_op.concat",
     "pd_op.split",
     "pd_op.multiply",
+    "pd_op.relu",
+    "pd_op.sigmoid",
 ]
 
 

test/prim/pir_prim/test_prim_sub_graph_backward_dynamic_shape.py

@@ -91,6 +91,14 @@ def multiply_net(x, y):
     return x * y
 
 
+def relu_net(x):
+    return paddle.nn.functional.relu(x)
+
+
+def sigmoid_net(x):
+    return paddle.nn.functional.sigmoid(x)
+
+
 def apply_to_static(net, use_cinn, input_spec=None):
     build_strategy = paddle.static.BuildStrategy()
     build_strategy.build_cinn_pass = use_cinn
@@ -782,5 +790,29 @@ def setUp(self):
         self.tol = 1e-5
 
 
+class TestPrimReluWithGrad(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 = relu_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimSigmoidWithGrad(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 = sigmoid_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
 if __name__ == "__main__":
     unittest.main()