Scalar support marking data_type in yaml (#40867)

* Scalar support marking data_type in yaml * fix code-gene bug
PaddlePaddle · Mar 25, 2022 · 0408701 · 0408701
1 parent 0f5e90a
commit 0408701
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 24 deletions.
diff --git a/paddle/fluid/eager/auto_code_generator/final_state_generator/codegen_utils.py b/paddle/fluid/eager/auto_code_generator/final_state_generator/codegen_utils.py
@@ -30,15 +30,19 @@
 core_ops_args_type_info = {}
 
 yaml_types_mapping = {
-    'int' : 'int', 'int32' : 'int32_t', 'int64' : 'int64_t',  'size_t' : 'size_t', \
+    'int' : 'int', 'int32_t' : 'int32_t', 'int64_t' : 'int64_t',  'size_t' : 'size_t', \
     'float' : 'float', 'double' : 'double', 'bool' : 'bool', \
     'str' : 'std::string', \
     'Place' : 'paddle::experimental::Place', 'DataLayout' : 'paddle::experimental::DataLayout', 'DataType' : 'paddle::experimental::DataType', \
-    'int64[]' : 'std::vector<int64_t>', 'int[]' : 'std::vector<int>',
+    'int64_t[]' : 'std::vector<int64_t>', 'int[]' : 'std::vector<int>',
     'Tensor' : 'Tensor',
     'Tensor[]' : 'std::vector<Tensor>',
     'Tensor[Tensor[]]' : 'std::vector<std::vector<Tensor>>',
     'Scalar' : 'paddle::experimental::Scalar',
+    'Scalar(int)' : 'paddle::experimental::Scalar',
+    'Scalar(int64_t)' : 'paddle::experimental::Scalar',
+    'Scalar(float)' : 'paddle::experimental::Scalar',
+    'Scalar(double)' : 'paddle::experimental::Scalar',
     'ScalarArray' : 'paddle::experimental::ScalarArray'
 }
 
@@ -254,8 +258,8 @@ def ParseYamlForward(args_str, returns_str):
 
     fargs = r'(.*?)'
     wspace = r'\s*'
-    args_pattern = f'\({fargs}\)'
-    args_str = re.search(args_pattern, args_str).group(1)
+    args_pattern = f'^\({fargs}\)$'
+    args_str = re.search(args_pattern, args_str.strip()).group(1)
 
     inputs_list, attrs_list = ParseYamlArgs(args_str)
     returns_list = ParseYamlReturns(returns_str)

diff --git a/python/paddle/utils/code_gen/api.yaml b/python/paddle/utils/code_gen/api.yaml
@@ -19,7 +19,7 @@
 
 
 - api : concat
-  args : (Tensor[] x, Scalar axis)
+  args : (Tensor[] x, Scalar(int64_t) axis)
   output : Tensor
   infer_meta :
     func : ConcatInferMeta
@@ -123,7 +123,7 @@
   backward : matmul_grad
 
 - api : mean
-  args : (Tensor x, int64[] axis={}, bool keep_dim=false)
+  args : (Tensor x, int64_t[] axis={}, bool keep_dim=false)
   output : Tensor
   infer_meta :
     func : ReduceInferMeta
@@ -198,7 +198,7 @@
     func : sotfmax
 
 - api : split
-  args : (Tensor x, ScalarArray num_or_sections, Scalar axis)
+  args : (Tensor x, ScalarArray num_or_sections, Scalar(int) axis)
   output : Tensor[]
   invoke : split_impl(x, num_or_sections, axis)
 
@@ -212,7 +212,7 @@
   backward : subtract_grad
 
 - api : sum
-  args : (Tensor x, int64[] axis={}, DataType dtype=DataType::UNDEFINED, bool keep_dim=false)
+  args : (Tensor x, int64_t[] axis={}, DataType dtype=DataType::UNDEFINED, bool keep_dim=false)
   output : Tensor
   infer_meta :
     func : SumInferMeta
@@ -227,7 +227,7 @@
 
 
 - api : one_hot
-  args : (Tensor x, Scalar num_classes)
+  args : (Tensor x, Scalar(int) num_classes)
   output : Tensor
   infer_meta :
     func : OneHotInferMeta

diff --git a/python/paddle/utils/code_gen/api_base.py b/python/paddle/utils/code_gen/api_base.py
@@ -89,10 +89,13 @@ def parse_input_and_attr(self, api_name, args_config, optional_vars=[]):
         attr_types_map = {
             'ScalarArray': 'const ScalarArray&',
             'Scalar': 'const Scalar&',
-            'uint8': 'uint8_t',
+            'Scalar(int)': 'const Scalar&',
+            'Scalar(int64_t)': 'const Scalar&',
+            'Scalar(float)': 'const Scalar&',
+            'Scalar(dobule)': 'const Scalar&',
             'int': 'int',
-            'int32': 'int32_t',
-            'int64': 'int64_t',
+            'int32_t': 'int32_t',
+            'int64_t': 'int64_t',
             'long': 'long',
             'size_t': 'size_t',
             'float': 'float',
@@ -102,27 +105,21 @@ def parse_input_and_attr(self, api_name, args_config, optional_vars=[]):
             'Place': 'Place',
             'DataLayout': 'DataLayout',
             'DataType': 'DataType',
-            'int64[]': 'const std::vector<int64_t>&',
-            'int[]': 'const std::vector<int>&',
-            'long[]': 'const std::vector<int64_t>&'
+            'int64_t[]': 'const std::vector<int64_t>&',
+            'int[]': 'const std::vector<int>&'
         }
         optional_types_trans = {
             'Tensor': 'const paddle::optional<Tensor>&',
             'Tensor[]': 'const paddle::optional<std::vector<Tensor>>&',
-            'ScalarArray': 'const paddle::optional<ScalarArray>&',
-            'Scalar': 'const paddle::optional<Scalar>&',
             'int': 'paddle::optional<int>',
-            'int32': 'paddle::optional<int32_t>',
-            'int64': 'paddle::optional<int64_t>',
-            'size_t': 'paddle::optional<size_t>',
+            'int32_t': 'paddle::optional<int32_t>',
+            'int64_t': 'paddle::optional<int64_t>',
             'float': 'paddle::optional<float>',
             'double': 'paddle::optional<double>',
             'bool': 'paddle::optional<bool>',
             'Place': 'paddle::optional<Place>',
             'DataLayout': 'paddle::optional<DataLayout>',
-            'DataType': 'paddle::optional<DataType>',
-            'int64[]': 'paddle::optional<std::vector<int64_t>>',
-            'int[]': 'paddle::optional<std::vector<int>>'
+            'DataType': 'paddle::optional<DataType>'
         }
 
         args_declare_str = ""

diff --git a/python/paddle/utils/code_gen/sparse_api.yaml b/python/paddle/utils/code_gen/sparse_api.yaml
@@ -13,7 +13,7 @@
   invoke : to_dense_impl(x)
 
 - api : to_sparse_coo
-  args : (Tensor x, int64 sparse_dim)
+  args : (Tensor x, int64_t sparse_dim)
   output : Tensor(out@SparseCooTensor)
   invoke : to_sparse_coo_impl(x, sparse_dim)