[Prim][PIR] Support the dynamic shape for dropout_grad (#67757)

zeroRains · web-flow · commit ae932e44ace6 · 2024-08-28T16:38:34.000+08:00
* surport dynamic shape for dropout grad

* fix the function position
diff --git a/paddle/fluid/primitive/rule/vjp/details.h b/paddle/fluid/primitive/rule/vjp/details.h
@@ -985,16 +985,30 @@ void dropout_grad(const Tensor& mask,
     if (mode == "upscale_in_train") {
       by_pass<T>(out_grad, x_grad);
     } else {
-      set_output<T>(out_grad * (1.0 - p.to<float>()), x_grad);
+      Tensor scalar = full_scalar<T>(1.0 - p.to<float>(), out_grad.dtype());
+      set_output<T>(out_grad * scalar, x_grad);
     }
   } else {
     if (mode == "upscale_in_train") {
-      if (p.to<float>() == 1.0f) {
-        set_output<T>(scale<T>(out_grad, 0.0), x_grad);
+      if (has_dynamic_shape(out_grad.shape())) {
+        if (p.to<float>() == 1.0f) {
+          Tensor zero = full_scalar<T>(0.0, out_grad.dtype());
+          set_output<T>(backend::scale<T>(out_grad, zero), x_grad);
+        } else {
+          Tensor scalar =
+              full_scalar<T>(1.0 / (1.0 - p.to<float>()), out_grad.dtype());
+          set_output<T>(backend::scale<T>(
+                            out_grad * cast<T>(mask, out_grad.dtype()), scalar),
+                        x_grad);
+        }
       } else {
-        set_output<T>(scale<T>(out_grad * cast<T>(mask, out_grad.dtype()),
-                               1.0 / (1.0 - p.to<float>())),
-                      x_grad);
+        if (p.to<float>() == 1.0f) {
+          set_output<T>(scale<T>(out_grad, 0.0), x_grad);
+        } else {
+          set_output<T>(scale<T>(out_grad * cast<T>(mask, out_grad.dtype()),
+                                 1.0 / (1.0 - p.to<float>())),
+                        x_grad);
+        }
       }
     } else {
       set_output<T>(out_grad * cast<T>(mask, out_grad.dtype()), x_grad);
diff --git a/python/paddle/autograd/backward_utils.py b/python/paddle/autograd/backward_utils.py
@@ -37,6 +37,7 @@
     "pd_op.cos",
     "pd_op.cumsum",
     "pd_op.divide",
+    "pd_op.dropout",
     "pd_op.elementwise_pow",
     "pd_op.erf",
     "pd_op.exp",
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
@@ -52,6 +52,20 @@ def divide_net(x, y):
     return x / y
 
 
+def dropout_net1(x):
+    return paddle.nn.functional.dropout(
+        x, training=False, mode='downscale_in_infer'
+    )
+
+
+def dropout_net2(x):
+    return paddle.nn.functional.dropout(x)
+
+
+def dropout_net3(x):
+    return paddle.nn.functional.dropout(x, 1.0)
+
+
 def elementwise_pow_net(x, y):
     return paddle.pow(x, y)
 
@@ -507,6 +521,56 @@ def setUp(self):
         self.tol = 1e-5
 
 
+class TestPrimDropoutWithGrad1(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 40, 50]
+        self.init_x_shape = [None, None, 50]
+        self.x = np.ones(self.x_shape).astype(self.dtype)
+        self.net = dropout_net1
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimDropoutWithGrad2(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [300, 4096]
+        self.init_x_shape = [None, 4096]
+        self.x = np.ones(self.x_shape).astype(self.dtype)
+        self.net = dropout_net2
+        self.enable_cinn = False
+        self.tol = 0.07
+
+    def test_prim_all_dynamic(self):
+        res_ref, grad_ref = self.base_net()
+        res, grad = self.base_net("prim")
+
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(
+                ref.sum(), actual.sum(), rtol=self.tol, atol=self.tol
+            )
+
+        for dr, d in zip(grad_ref, grad):
+            np.testing.assert_allclose(
+                dr.sum(), d.sum(), rtol=self.tol, atol=self.tol
+            )
+
+
+class TestPrimDropoutWithGrad3(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 40, 50]
+        self.init_x_shape = [None, None, 50]
+        self.x = np.ones(self.x_shape).astype(self.dtype)
+        self.net = dropout_net3
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
 class TestPrimElementwisePowWithGrad1(TestPrimTwoWithGrad):
     def setUp(self):
         np.random.seed(2023)
