[Dy2stat] Refactoring tensor_shape_transformer.py to Fix Change after Assign Bug (#31082)

**Problem**
In our old shape transformer logic, if user write:
```
s = tensor.shape
...
y = paddle.some_api(s)
```
Dy2stat will change it to
```
...
y = paddle.some_api(convert_var_shape(tensor))
```
However it will cause fatal bug if user changes the shape of `x` after assign. For example:
```
s = tensor.shape
...
tensor = paddle.some_change_shape_api(tensor)
...
y = paddle.some_api(s)
```
Then the Dy2stat will get wrong result because the code is translated into:
```
tensor = paddle.some_change_shape_api(tensor)
...
y = paddle.some_api(convert_var_shape(tensor)) # tensor shape has been changed, not origin `s` value
```

**Solution Logic**

It can not be solved in the old logic, so I refactoring tensor_shape_transformer logic. Now we will use `s` to store shape attribute and generate a var `s__STATIC_CONVERT_VAR_SHAPE_SUFFIX` to store static shape API `shape(tensor)`
```
s = tensor.shape
...
y = paddle.some_api(s)
```
Dy2stat will change it to
```
s = tensor.shape
s__STATIC_CONVERT_VAR_SHAPE_SUFFIX = shape(tensor)
...
y = paddle.some_api(choose_shape_attr_or_api(s, s__STATIC_CONVERT_VAR_SHAPE_SUFFIX ))
```
In this case, the code is consistent with origin dygraph meaning and it fixed the change after assign bug.

**Code Key Note**

To help reviewers, the key change of this PR is changing `self.name_to_var_shape` from "mapping name to shape node" to "mapping name to its STATIC_CONVERT_VAR_SHAPE_SUFFIX name", then if a variable name has the SUFFIX, we can choose to use attribute shape or shape api. Other changes go with the key change.

**Consideration**
The issue of this PR is that we store extra static `shape` API result, will it harms the speed of Dy2stat? In some cases it will, but we argue that the benefit would be greater than the cost.

1. The extra calling to static `shape` API will happen when coder assign among shape variables. Take the following dygraph code as an instance:
```
s1 = tensor.shape
s2 = s1
s3 = s2
...
```
Then we called extra static `shape` APIs again and again, however users seldom write code like this.

2. If the shape variable is used a lot, for example:
```
s = tensor.shape
y1 = paddle.some_api1(s)
y2 = paddle.some_api2(s)
y3 = paddle.some_api3(s)
```
Our old logic will create 3 shape APIs but now just 1. This is more common user code pattern. In fact, if reviewers take a look at the current unit test in this PR, you could see the op numbers decrease after this PR. So we argue that this PR can also improve speed in this code pattern.

python/paddle/fluid/dygraph/dygraph_to_static/convert_operators.py

@@ -267,12 +267,12 @@ def convert_var_shape(x, idx=None, in_control_flow=False):
     A function representation of the shape of variable.
     """
 
-    def has_negetive(list_shape, idx=None):
+    def has_negative(list_shape, idx=None):
         if idx is not None:
             return list_shape[idx] < 0
 
-        num_negetive = sum([1 if i < 0 else 0 for i in list_shape])
-        return num_negetive > 0
+        num_negative = sum([1 if i < 0 else 0 for i in list_shape])
+        return num_negative > 0
 
     # When `x` is Variable, call nn.shape(x) in following cases:
     #  (1) The shape of `x` is used in control flow condition.
@@ -280,18 +280,62 @@ def convert_var_shape(x, idx=None, in_control_flow=False):
     #      if x.shape[0] == 1:
     #          y = XX
     #      ```
-    #  (2) The dim to be used is negetive
+    #  (2) The dim to be used is negative
     #      ```
     #      # Assume x.shape=[3, -1] in static mode
     #      y = paddle.reshape(x, shape=[1, x.shape[1]])
     #      ```
-    if isinstance(x, Variable) and (in_control_flow or has_negetive(x.shape,
+    if isinstance(x, Variable) and (in_control_flow or has_negative(x.shape,
                                                                     idx)):
         return nn.shape(x) if idx is None else nn.shape(x)[idx]
     else:
         return x.shape if idx is None else x.shape[idx]
 
 
+def convert_var_shape_simple(x):
+    """
+    A function representation of the shape of variable.
+    """
+    if isinstance(x, Variable):
+        return nn.shape(x)
+    else:
+        return x.shape
+
+
+def eval_if_exist_else_none(name):
+    try:
+        return eval(name)
+    except:
+        return None
+
+
+def choose_shape_attr_or_api(attr_shape, api_shape, idx=None):
+    """
+    Input can be attribute `x.shape` or api `shape(x)`, this function
+    chooses which one to return to use in dy2stat.
+
+    Note: sometimes users write `x.shape[3]`, so attr_shape can be an integer.
+    """
+    if api_shape is None:
+        return attr_shape if idx is None else attr_shape[idx]
+    if not isinstance(attr_shape, (list, tuple)):
+        # some variables like x.shape[0] is no longer a list or tuple
+        if isinstance(attr_shape, int) and attr_shape < 0:
+            return api_shape if idx is None else api_shape[idx]
+        return attr_shape if idx is None else attr_shape[idx]
+
+    def has_negative(list_shape, idx=None):
+        if idx is not None:
+            return list_shape[idx] < 0
+
+        num_negative = sum([1 if i < 0 else 0 for i in list_shape])
+        return num_negative > 0
+
+    if has_negative(attr_shape, idx):
+        return api_shape if idx is None else api_shape[idx]
+    return attr_shape if idx is None else attr_shape[idx]
+
+
 def convert_shape_compare(left, *args):
     """
     A function handles comparison difference between Paddle and Python.

python/paddle/fluid/tests/unittests/dygraph_to_static/test_convert_operators.py

@@ -136,5 +136,58 @@ class TestConvertShapeCompare(unittest.TestCase):
         paddle.disable_static()
 
 
+class TestChooseShapeAttrOrApi(unittest.TestCase):
+    def test_api_shape_is_none(self):
+        self.assertEqual(
+            paddle.jit.dy2static.choose_shape_attr_or_api([1, 2], None),
+            [1, 2])
+        self.assertEqual(
+            paddle.jit.dy2static.choose_shape_attr_or_api([1], None), [1])
+        self.assertEqual(
+            paddle.jit.dy2static.choose_shape_attr_or_api([2, 3, 7], None, 0),
+            2)
+
+    def test_attr_shape_is_int(self):
+        x = paddle.zeros([1, 3, 5, 7])
+        self.assertEqual(
+            paddle.jit.dy2static.choose_shape_attr_or_api(x.shape[0],
+                                                          paddle.shape(x)[0]),
+            1)
+        self.assertEqual(
+            paddle.jit.dy2static.choose_shape_attr_or_api(x.shape[1],
+                                                          paddle.shape(x)[1]),
+            3)
+        self.assertEqual(
+            paddle.jit.dy2static.choose_shape_attr_or_api(-1,
+                                                          paddle.shape(x)[0]),
+            paddle.shape(x)[0])
+        self.assertEqual(
+            paddle.jit.dy2static.choose_shape_attr_or_api(-1,
+                                                          paddle.shape(x), 0),
+            paddle.shape(x)[0])
+
+    def test_positive_attr_shape(self):
+        x = paddle.zeros([1, 3, 5, 7])
+        self.assertEqual(
+            paddle.jit.dy2static.choose_shape_attr_or_api(x.shape,
+                                                          paddle.shape(x)),
+            x.shape)
+        self.assertEqual(
+            paddle.jit.dy2static.choose_shape_attr_or_api(x.shape,
+                                                          paddle.shape(x), 3),
+            x.shape[3])
+
+    def test_negative_attr_shape(self):
+        x = paddle.zeros([7])
+        self.assertEqual(
+            paddle.jit.dy2static.choose_shape_attr_or_api([-1],
+                                                          paddle.shape(x), 0),
+            paddle.shape(x)[0])
+        self.assertEqual(
+            paddle.jit.dy2static.choose_shape_attr_or_api([-1],
+                                                          paddle.shape(x)),
+            paddle.shape(x))
+
+
 if __name__ == '__main__':
     unittest.main()

python/paddle/fluid/tests/unittests/dygraph_to_static/test_tensor_shape.py

@@ -207,6 +207,14 @@ def dyfunc_with_while_4(x):
     return x
 
 
+def dyfunc_change_shape_after_assign(x):
+    x = paddle.to_tensor(x)
+    a, b = x.shape
+    x = paddle.reshape(x, shape=(-1, 1))
+    res = paddle.reshape(x, shape=(b, a))
+    return res
+
+
 # 1. Basic tests without control flow
 class TestTensorShapeBasic(unittest.TestCase):
     def setUp(self):
@@ -289,11 +297,21 @@ class TestTensorShapeBasic5(TestTensorShapeBasic):
     def init_test_func(self):
         self.dygraph_func = dyfunc_tensor_shape_5
 
+    def _set_expected_op_num(self):
+        self.expected_op_num = 4
+        self.expected_shape_op_num = 1
+        self.expected_slice_op_num = 1
+
 
 class TestTensorShapeBasic6(TestTensorShapeBasic):
     def init_test_func(self):
         self.dygraph_func = dyfunc_tensor_shape_6
 
+    def _set_expected_op_num(self):
+        self.expected_op_num = 4
+        self.expected_shape_op_num = 1
+        self.expected_slice_op_num = 1
+
 
 class TestTupleShape1(TestTensorShapeBasic):
     def init_test_func(self):
@@ -327,9 +345,9 @@ class TestTensorShapeInIf1(TestTensorShapeBasic):
         self.dygraph_func = dyfunc_with_if_1
 
     def _set_expected_op_num(self):
-        self.expected_op_num = 26
-        self.expected_shape_op_num = 2
-        self.expected_slice_op_num = 2
+        self.expected_op_num = 4
+        self.expected_shape_op_num = 1
+        self.expected_slice_op_num = 1
 
 
 class TestTensorShapeInIf2(TestTensorShapeBasic):
@@ -357,6 +375,11 @@ class TestTensorShapeInFor2(TestTensorShapeInFor1):
     def init_test_func(self):
         self.dygraph_func = dyfunc_with_for_2
 
+    def _set_expected_op_num(self):
+        self.expected_op_num = 9
+        self.expected_shape_op_num = 1
+        self.expected_slice_op_num = 1
+
 
 # 4. Tests with control flow while loop
 class TestTensorShapeInWhile1(TestTensorShapeInFor1):
@@ -368,15 +391,20 @@ class TestTensorShapeInWhile2(TestTensorShapeInFor1):
     def init_test_func(self):
         self.dygraph_func = dyfunc_with_while_2
 
+    def _set_expected_op_num(self):
+        self.expected_op_num = 6
+        self.expected_shape_op_num = 1
+        self.expected_slice_op_num = 1
+
 
 class TestTensorShapeInWhile3(TestTensorShapeBasic):
     def init_test_func(self):
         self.dygraph_func = dyfunc_with_while_3
 
     def _set_expected_op_num(self):
-        self.expected_op_num = 25
-        self.expected_shape_op_num = 6
-        self.expected_slice_op_num = 3
+        self.expected_op_num = 2
+        self.expected_shape_op_num = 0
+        self.expected_slice_op_num = 0
 
 
 class TestTensorShapeInWhile4(TestTensorShapeBasic):
@@ -446,9 +474,9 @@ class TestOpNumWithTensorShapeTuple1(TestOpNumBasicWithTensorShape):
         self.dygraph_func = dyfunc_tuple_shape_1
 
     def _set_expected_op_num(self):
-        self.expected_op_num = 5
-        self.expected_shape_op_num = 1
-        self.expected_slice_op_num = 1
+        self.expected_op_num = 2
+        self.expected_shape_op_num = 0
+        self.expected_slice_op_num = 0
 
 
 class TestOpNumWithTensorShapeInIf1(TestOpNumBasicWithTensorShape):
@@ -456,7 +484,7 @@ class TestOpNumWithTensorShapeInIf1(TestOpNumBasicWithTensorShape):
         self.dygraph_func = dyfunc_with_if_1
 
     def _set_expected_op_num(self):
-        self.expected_op_num = 28
+        self.expected_op_num = 19
         self.expected_shape_op_num = 4
         self.expected_slice_op_num = 2
 
@@ -481,5 +509,17 @@ class TestOpNumWithTensorShapeInWhile1(TestOpNumBasicWithTensorShape):
         self.expected_slice_op_num = 3
 
 
+class TestChangeShapeAfterAssign(TestTensorShapeBasic):
+    def init_test_func(self):
+        self.input = numpy.ones((2, 3)).astype("int32")
+        self.input_spec = [paddle.static.InputSpec(shape=[2, 3], dtype="int32")]
+        self.dygraph_func = dyfunc_change_shape_after_assign
+
+    def _set_expected_op_num(self):
+        self.expected_op_num = 3
+        self.expected_shape_op_num = 0
+        self.expected_slice_op_num = 0
+
+
 if __name__ == '__main__':
     unittest.main()

python/paddle/jit/dy2static/convert_operators.py

@@ -25,11 +25,15 @@ from ...fluid.dygraph.dygraph_to_static.convert_operators import convert_print
 from ...fluid.dygraph.dygraph_to_static.convert_operators import convert_shape_compare  #DEFINE_ALIAS
 from ...fluid.dygraph.dygraph_to_static.convert_operators import convert_var_dtype  #DEFINE_ALIAS
 from ...fluid.dygraph.dygraph_to_static.convert_operators import convert_var_shape  #DEFINE_ALIAS
+from ...fluid.dygraph.dygraph_to_static.convert_operators import convert_var_shape_simple  #DEFINE_ALIAS
+from ...fluid.dygraph.dygraph_to_static.convert_operators import eval_if_exist_else_none  #DEFINE_ALIAS
+from ...fluid.dygraph.dygraph_to_static.convert_operators import choose_shape_attr_or_api  #DEFINE_ALIAS
 from ...fluid.dygraph.dygraph_to_static.convert_operators import convert_while_loop  #DEFINE_ALIAS
 
 __all__ = [
     'cast_bool_if_necessary', 'convert_assert', 'convert_ifelse', 'convert_len',
     'convert_logical_and', 'convert_logical_not', 'convert_logical_or',
     'convert_pop', 'convert_print', 'convert_shape_compare',
-    'convert_var_dtype', 'convert_var_shape', 'convert_while_loop'
+    'convert_var_dtype', 'convert_var_shape', 'convert_var_shape_simple',
+    'eval_if_exist_else_none', 'choose_shape_attr_or_api', 'convert_while_loop'
 ]