Add sqrt orig2prim rule and UT

python/paddle/autograd/primrules.py

@@ -112,6 +112,7 @@ def linear_jvp(op, *args, **kwargs):
 scale
 assign
 elementwise_mul
+sqrt
 
 These original ops are partially supported:
 
@@ -123,6 +124,11 @@ def linear_jvp(op, *args, **kwargs):
 """
 
 
+@REGISTER_ORIG2PRIM('sqrt')
+def sqrt_orig2prim(op, x):
+    return sqrt(x)
+
+
 @REGISTER_ORIG2PRIM('elementwise_mul')
 def elementwise_mul_orig2prim(op, x, y):
     if x.shape != y.shape:

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

@@ -68,6 +68,24 @@ def test_op(self):
                 self.assertEqual(prim_out[k].shape, v.shape)
 
 
+class TestSqrtOrig2Prim(TestElementWiseAddOrig2Prim):
+    def init_data(self):
+        self.op_type = 'sqrt'
+        X = paddle.static.data(name='X', shape=[7, 8], dtype='float64')
+
+        self.input = {'X': X, }
+        self.output = {
+            'Out':
+            self.layer_help.create_variable_for_type_inference(dtype=X.dtype)
+        }
+        self.attrs = {}
+
+        self.orig2prim_args = (X, )
+        self.all_ops = ['sqrt', 'sqrt_p']
+        # { prim_op_output_index: orig_op_output_var }
+        self.out_map = {0: self.output['Out']}
+
+
 class TestElementWiseMulOrig2Prim(TestElementWiseAddOrig2Prim):
     def init_data(self):
         self.op_type = 'elementwise_mul'