[Dy2St]Handle a, b = paddle.shape(x) in Static Analysis (#39245)

* refine Assign * add UT
PaddlePaddle · Feb 10, 2022 · 1252f4b · 1252f4b
1 parent 383de29
commit 1252f4b
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/python/paddle/fluid/dygraph/dygraph_to_static/static_analysis.py b/python/paddle/fluid/dygraph/dygraph_to_static/static_analysis.py
@@ -334,6 +334,13 @@ def _get_node_var_type(self, cur_wrapper):
                 if isinstance(target, gast.Name):
                     self.node_to_wrapper_map[target].node_var_type = ret_type
                     self.var_env.set_var_type(target.id, ret_type)
+                # Handle statements like `a, b = paddle.shape(x)`
+                elif isinstance(target, gast.Tuple):
+                    for sub_target in target.elts:
+                        if isinstance(sub_target, gast.Name):
+                            self.node_to_wrapper_map[
+                                sub_target].node_var_type = ret_type
+                            self.var_env.set_var_type(sub_target.id, ret_type)
             return ret_type
 
         if isinstance(node, gast.AnnAssign):

diff --git a/python/paddle/fluid/tests/unittests/dygraph_to_static/test_static_analysis.py b/python/paddle/fluid/tests/unittests/dygraph_to_static/test_static_analysis.py
@@ -146,13 +146,16 @@ def add(x, y):
 
 def func_to_test7(a: int, b: float, c: paddle.Tensor, d: float='diff'):
     a = True
+    e, f = paddle.shape(c)
 
 
 result_var_type7 = {
     'a': {NodeVarType.BOOLEAN},
     'b': {NodeVarType.FLOAT},
     'c': {NodeVarType.TENSOR},
-    'd': {NodeVarType.STRING}
+    'd': {NodeVarType.STRING},
+    'e': {NodeVarType.PADDLE_RETURN_TYPES},
+    'f': {NodeVarType.PADDLE_RETURN_TYPES}
 }
 
 test_funcs = [