Fix dropout static when axis != None

python/paddle/fluid/tests/unittests/test_dropout_op.py

@@ -333,7 +333,7 @@ def setUp(self):
 
     def check_static_result(self, place):
         with fluid.program_guard(fluid.Program(), fluid.Program()):
-            input = fluid.data(name="input", shape=[40, 40], dtype="float32")
+            input = fluid.data(name="input", shape=[-1, -1], dtype="float32")
             res1 = paddle.nn.functional.dropout(x=input, p=0., training=False)
             res2 = paddle.nn.functional.dropout(
                 x=input, p=0., axis=0, training=True, mode='upscale_in_train')
@@ -380,7 +380,11 @@ def check_static_result(self, place):
                 training=False,
                 mode='upscale_in_train')
 
-            in_np = np.random.random([40, 40]).astype("float32")
+            res13 = paddle.nn.functional.dropout(
+                x=input, p=0.7, axis=1, training=True, mode='upscale_in_train')
+
+            in_np = np.ones([40, 40]).astype("float32")
+            in_np2 = np.ones([1, 250000000]).astype("float32")
             res_np = in_np
             res_np2 = np.zeros_like(in_np)
 
@@ -398,6 +402,12 @@ def check_static_result(self, place):
                                feed={"input": in_np},
                                fetch_list=[res10])
             self.assertTrue(np.allclose(fetches2[0], res_np2))
+            fetches3 = exe.run(fluid.default_main_program(),
+                               feed={"input": in_np2},
+                               fetch_list=[res13])
+            self.assertTrue(
+                np.isclose(
+                    np.sum(fetches3[0] == 0) / np.sum(in_np2), 0.7, atol=1e-04))
 
     def test_static(self):
         for place in self.places:

python/paddle/nn/functional/common.py

@@ -940,6 +940,8 @@ def get_attrs(prog, dropout_prob, is_test, seed):
 
             #get mask shape
             input_shape = x.shape
+            if not in_dygraph_mode():
+                input_shape_tensor = paddle.shape(x)
             drop_axes = [axis] if isinstance(axis, int) else list(axis)
             if min(drop_axes) < 0 or max(drop_axes) > len(input_shape) - 1:
                 raise ValueError("axis value should be greater than or equal to 0 and less than dimensions of x:{}, but get axis value:{} " \
@@ -949,8 +951,12 @@ def get_attrs(prog, dropout_prob, is_test, seed):
                     "length of axis should not be greater than dimensions of x:{}, but get length of axis: {}".
                     format(len(input_shape), len(drop_axes)))
             mask_shape = [1] * len(input_shape)
-            for i in drop_axes:
-                mask_shape[i] = input_shape[i]
+            if not in_dygraph_mode():
+                for i in drop_axes:
+                    mask_shape[i] = input_shape_tensor[i]
+            else:
+                for i in drop_axes:
+                    mask_shape[i] = input_shape[i]
 
             #get mask
             random_tensor = paddle.uniform(