add gelu orig2prim rule

python/paddle/fluid/tests/unittests/autograd/test_orig2prim.py

@@ -579,5 +579,50 @@ def init_data(self):
         self.out_map = {0: self.output['Out']}
 
 
+class TestGeluOrig2Prim(TestElementWiseAddOrig2Prim):
+
+    def init_data(self):
+        self.op_type = 'gelu'
+        X = paddle.static.data(name='X', shape=[5, 8], dtype='float')
+
+        self.input = {'X': X}
+        self.output = {
+            'Out':
+            self.layer_help.create_variable_for_type_inference(dtype=X.dtype)
+        }
+        self.attrs = {'approximate': False}
+
+        self.orig2prim_args = (X, )
+        self.all_ops = [
+            'gelu', 'add_p', 'erf_p', 'fill_constant_p', 'fill_constant_p',
+            'fill_constant_p', 'mul_p', 'mul_p', 'mul_p'
+        ]
+        # { prim_op_output_index: orig_op_output_var }
+        self.out_map = {0: self.output['Out']}
+
+
+class TestGeluApproximateOrig2Prim(TestElementWiseAddOrig2Prim):
+
+    def init_data(self):
+        self.op_type = 'gelu'
+        X = paddle.static.data(name='X', shape=[5, 8], dtype='float')
+
+        self.input = {'X': X}
+        self.output = {
+            'Out':
+            self.layer_help.create_variable_for_type_inference(dtype=X.dtype)
+        }
+        self.attrs = {'approximate': True}
+
+        self.orig2prim_args = (X, )
+        self.all_ops = [
+            'add_p', 'add_p', 'fill_constant_p', 'fill_constant_p',
+            'fill_constant_p', 'fill_constant_p', 'fill_constant_p', 'gelu',
+            'mul_p', 'mul_p', 'mul_p', 'mul_p', 'pow_p', 'tanh_p'
+        ]
+        # { prim_op_output_index: orig_op_output_var }
+        self.out_map = {0: self.output['Out']}
+
+
 if __name__ == '__main__':
     unittest.main()

python/paddle/fluid/tests/unittests/autograd/test_primapi.py

@@ -279,6 +279,10 @@ def test_illegal_param(self):
             np.array([2, 2, 2]),
         ), None, 'float32'),
         ('erf', paddle.erf, (np.random.rand(300, 288), ), None, 'float32'),
+        ('gelu', paddle.nn.functional.gelu,
+         (np.random.rand(200, 189), ), None, 'float32'),
+        ('gelu_approximate', lambda x: paddle.nn.functional.gelu(x, True),
+         (np.random.rand(200, 189), ), None, 'float32'),
     ))
 class TestGrad(unittest.TestCase):
 
@@ -406,20 +410,33 @@ def multiply_pd(x):
 erf_ag = lambda xs: ascipy.special.erf(xs[0])
 
 
+def gelu_ag(x, approximate=False):
+    if approximate:
+        sqrt_2_over_pi = np.sqrt(2 / np.pi).astype(x.dtype)
+        cdf = 0.5 * (1.0 + anp.tanh(sqrt_2_over_pi * (x + 0.044715 * (x**3))))
+        return x * cdf
+    else:
+        return x * (ascipy.special.erf(x / np.sqrt(2)) + 1) / 2
+
+
 @utils.place(config.DEVICES)
 @utils.parameterize(
-    (utils.TEST_CASE_NAME, 'fun_pd', 'fun_ag', 'xs', 'v', 'dtype'), (
-        ('multiply', multiply_pd, multiply_ag,
-         (np.random.rand(3, 5), ), None, 'float32'),
-        ('sin', paddle.sin, sin_ag, (np.random.rand(2, 3), ), None, 'float32'),
-        ('cos', paddle.cos, cos_ag, (np.random.rand(3, 4), ), None, 'float32'),
-        ('exp', paddle.exp, exp_ag, (np.random.rand(2, 3), ), None, 'float32'),
-        ('pow', paddle.pow, pow_ag,
-         (np.random.rand(2, 3), np.random.rand(2, 3)), None, 'float32'),
-        ('log', paddle.log, log_ag, (np.random.rand(3, 8), ), None, 'float32'),
-        ('erf', paddle.erf, erf_ag,
-         (np.random.rand(100, 200), ), None, 'float32'),
-    ))
+    (utils.TEST_CASE_NAME, 'fun_pd', 'fun_ag', 'xs', 'v', 'dtype'),
+    (('multiply', multiply_pd, multiply_ag,
+      (np.random.rand(3, 5), ), None, 'float32'),
+     ('sin', paddle.sin, sin_ag, (np.random.rand(2, 3), ), None, 'float32'),
+     ('cos', paddle.cos, cos_ag, (np.random.rand(3, 4), ), None, 'float32'),
+     ('exp', paddle.exp, exp_ag, (np.random.rand(2, 3), ), None, 'float32'),
+     ('pow', paddle.pow, pow_ag,
+      (np.random.rand(2, 3), np.random.rand(2, 3)), None, 'float32'),
+     ('log', paddle.log, log_ag, (np.random.rand(3, 8), ), None, 'float32'),
+     ('erf', paddle.erf, erf_ag, (np.random.rand(100, 200), ), None, 'float32'),
+     ('gelu', paddle.nn.functional.gelu, lambda xs: gelu_ag(xs[0]),
+      (np.random.rand(10, 20, 30), ), None, 'float32'),
+     ('gelu_approximate',
+      lambda x: paddle.nn.functional.gelu(x, approximate=True),
+      lambda xs: gelu_ag(xs[0], approximate=True),
+      (np.random.rand(10, 20, 30), ), None, 'float32')))
 class TestGradWithHigherOrder(unittest.TestCase):
 
     def setUp(self):

python/paddle/incubate/autograd/primrules.py

@@ -327,13 +327,34 @@ def elementwise_pow_orig2prim(op, x, y):
 def elementwise_max_orig2prim(op, x, y):
     if x.shape != y.shape:
         y = broadcast(y, shape=x.shape)
-
     return primops.max(x, y)
 
 
-## Register prim2orig lower rules
+@REGISTER_ORIG2PRIM('gelu')
+def gelu_orig2prim(op, x):
+    if op.attr('approximate'):
+        cdf = mul(
+            fill_const(0.5, x.shape, x.dtype),
+            add(
+                fill_const(1.0, x.shape, x.dtype),
+                tanh(
+                    mul(
+                        fill_const(math.sqrt(2 / math.pi), x.shape, x.dtype),
+                        add(
+                            x,
+                            mul(
+                                fill_const(0.044715, x.shape, x.dtype),
+                                primops.pow(x, fill_const(3., x.shape,
+                                                          x.dtype))))))))
+        return mul(x, cdf)
+    else:
+        return mul(
+            mul(fill_const(0.5, x.shape, x.dtype), x),
+            add(fill_const(1.0, x.shape, x.dtype),
+                erf(mul(x, fill_const(1 / math.sqrt(2.), x.shape, x.dtype)))))
 
 
+## Register prim2orig lower rules
 @REGISTER_PRIM2ORIG('add_p')
 def add_prim2orig(op, x, y):
     return paddle.add(x, y)