[PIR] Fix 5 uts in PIR mode (test_imperative_mnist, test_imperative_mnist_sorted_gradient, test_imperative_reinforcement, test_imperative_transformer_sorted_gradient, test_inverse_op) (#65759)

SigureMo · web-flow · commit dabf0e034591 · 2024-07-08T10:28:38.000+08:00
diff --git a/python/paddle/static/input.py b/python/paddle/static/input.py
@@ -142,6 +142,10 @@ def _reset_data_op_insertion_point():
     for i in range(len(shape)):
         if shape[i] is None:
             shape[i] = -1
+        if isinstance(shape[i], int) and shape[i] < 0 and shape[i] != -1:
+            raise ValueError(
+                f"Only -1 can be used in shape to indicate unknown dimension, but received {shape[i]}"
+            )
 
     if dtype is None:
         dtype = paddle.get_default_dtype()
diff --git a/test/deprecated/legacy_test/CMakeLists.txt b/test/deprecated/legacy_test/CMakeLists.txt
@@ -394,8 +394,6 @@ endfunction()
 list(REMOVE_ITEM TEST_OPS test_feed_data_check_shape_type_deprecated)
 list(REMOVE_ITEM TEST_OPS test_fetch_lod_tensor_array)
 list(REMOVE_ITEM TEST_OPS test_data_norm_op_deprecated)
-list(REMOVE_ITEM TEST_OPS test_imperative_mnist_sorted_gradient)
-list(REMOVE_ITEM TEST_OPS test_imperative_mnist)
 list(REMOVE_ITEM TEST_OPS test_layers_deprecated)
 list(REMOVE_ITEM TEST_OPS test_imperative_ocr_attention_model)
 list(REMOVE_ITEM TEST_OPS test_install_check)
@@ -474,11 +472,6 @@ foreach(TEST_OP ${TEST_OPS})
 endforeach()
 set_tests_properties(test_logcumsumexp_op PROPERTIES TIMEOUT 30)
 
-py_test_modules(test_imperative_mnist MODULES test_imperative_mnist ENVS
-                FLAGS_cudnn_deterministic=1)
-py_test_modules(
-  test_imperative_mnist_sorted_gradient MODULES
-  test_imperative_mnist_sorted_gradient ENVS FLAGS_cudnn_deterministic=1)
 py_test_modules(
   test_imperative_ocr_attention_model MODULES
   test_imperative_ocr_attention_model ENVS FLAGS_cudnn_deterministic=1)
@@ -645,8 +638,6 @@ endif()
 set_tests_properties(test_add_reader_dependency_deprecated PROPERTIES TIMEOUT
                                                                       120)
 set_tests_properties(test_fleet_util PROPERTIES TIMEOUT 120)
-set_tests_properties(test_imperative_transformer_sorted_gradient
-                     PROPERTIES TIMEOUT 120)
 
 if(WITH_NV_JETSON)
   set_tests_properties(test_conv3d_transpose_part2_op_deprecated
@@ -681,7 +672,6 @@ set_tests_properties(test_fuse_bn_act_pass_deprecated PROPERTIES TIMEOUT 120)
 set_tests_properties(test_conv2d_api_deprecated PROPERTIES TIMEOUT 120)
 set_tests_properties(test_dygraph_multi_forward PROPERTIES TIMEOUT 120)
 set_tests_properties(test_imperative_ocr_attention_model PROPERTIES TIMEOUT 120)
-set_tests_properties(test_imperative_mnist PROPERTIES TIMEOUT 120)
 set_tests_properties(test_regularizer_deprecated PROPERTIES TIMEOUT 150)
 set_tests_properties(test_slice_op PROPERTIES TIMEOUT 120)
 set_tests_properties(test_slice_op_deprecated PROPERTIES TIMEOUT 120)
@@ -701,11 +691,6 @@ if(APPLE)
   set_tests_properties(test_callback_early_stop PROPERTIES TIMEOUT 300)
 endif()
 
-if(APPLE)
-  set_tests_properties(test_imperative_transformer_sorted_gradient
-                       PROPERTIES TIMEOUT 300)
-endif()
-
 set_tests_properties(test_inplace_addto_strategy_deprecated PROPERTIES TIMEOUT
                                                                        120)
 
diff --git a/test/ir/pir/test_data_op.py b/test/ir/pir/test_data_op.py
@@ -49,6 +49,20 @@ def test_with_pir(self):
                 out = data()
 
 
+class TestDataOpError(unittest.TestCase):
+    def test_invalid_shape(self):
+        paddle.enable_static()
+        place = paddle.CPUPlace()
+        exe = paddle.static.Executor(place)
+
+        main_program = paddle.static.Program()
+        new_scope = paddle.static.Scope()
+        with paddle.static.scope_guard(new_scope):
+            with paddle.static.program_guard(main_program):
+                with self.assertRaises(ValueError):
+                    out = paddle.static.data("x", shape=[-2, 1])
+
+
 if __name__ == "__main__":
     paddle.enable_static()
     unittest.main()
diff --git a/test/legacy_test/CMakeLists.txt b/test/legacy_test/CMakeLists.txt
@@ -429,6 +429,8 @@ endfunction()
 
 list(REMOVE_ITEM TEST_OPS test_data_norm_op)
 list(REMOVE_ITEM TEST_OPS test_warpctc_op)
+list(REMOVE_ITEM TEST_OPS test_imperative_mnist_sorted_gradient)
+list(REMOVE_ITEM TEST_OPS test_imperative_mnist)
 list(REMOVE_ITEM TEST_OPS test_imperative_resnet)
 list(REMOVE_ITEM TEST_OPS test_imperative_resnet_sorted_gradient)
 list(REMOVE_ITEM TEST_OPS test_imperative_se_resnext)
@@ -539,15 +541,21 @@ if(WITH_GPU
 endif()
 py_test_modules(test_imperative_resnet MODULES test_imperative_resnet ENVS
                 FLAGS_cudnn_deterministic=1)
-set_tests_properties(test_imperative_resnet
-                     PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE:NIGHTLY")
+py_test_modules(test_imperative_mnist MODULES test_imperative_mnist ENVS
+                FLAGS_cudnn_deterministic=1)
+py_test_modules(
+  test_imperative_mnist_sorted_gradient MODULES
+  test_imperative_mnist_sorted_gradient ENVS FLAGS_cudnn_deterministic=1)
 py_test_modules(
   test_imperative_resnet_sorted_gradient MODULES
   test_imperative_resnet_sorted_gradient ENVS FLAGS_cudnn_deterministic=1)
-set_tests_properties(test_imperative_resnet_sorted_gradient
-                     PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE:NIGHTLY")
 py_test_modules(test_imperative_se_resnext MODULES test_imperative_se_resnext
                 ENVS FLAGS_cudnn_deterministic=1)
+set_tests_properties(test_imperative_mnist PROPERTIES TIMEOUT 120)
+set_tests_properties(test_imperative_resnet
+                     PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE:NIGHTLY")
+set_tests_properties(test_imperative_resnet_sorted_gradient
+                     PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE:NIGHTLY")
 set_tests_properties(test_imperative_se_resnext
                      PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE:NIGHTLY")
 
@@ -861,6 +869,8 @@ set_tests_properties(
   PROPERTIES TIMEOUT 120)
 set_tests_properties(test_pool_max_op PROPERTIES TIMEOUT 500)
 set_tests_properties(test_group_norm_op PROPERTIES TIMEOUT 1000)
+set_tests_properties(test_imperative_transformer_sorted_gradient
+                     PROPERTIES TIMEOUT 120)
 set_tests_properties(test_imperative_optimizer PROPERTIES TIMEOUT 250)
 set_tests_properties(test_activation_op PROPERTIES TIMEOUT 270)
 set_tests_properties(test_normal PROPERTIES TIMEOUT 120)
@@ -900,6 +910,8 @@ if(APPLE)
   set_tests_properties(test_callback_reduce_lr_on_plateau PROPERTIES TIMEOUT
                                                                      300)
   set_tests_properties(test_vision_models PROPERTIES TIMEOUT 300)
+  set_tests_properties(test_imperative_transformer_sorted_gradient
+                       PROPERTIES TIMEOUT 300)
 endif()
 set_tests_properties(test_decorator PROPERTIES TIMEOUT 120)
 
diff --git a/test/legacy_test/test_imperative_mnist.py b/test/legacy_test/test_imperative_mnist.py
@@ -20,10 +20,35 @@
 
 import paddle
 from paddle import base
+from paddle.autograd.backward_utils import ValueDict
 from paddle.base import core
 from paddle.nn import Linear
 
 
+def create_parameter_mapping(startup_program, main_program):
+    startup_params = {}
+    main_params = {}
+    parameter_mapping = ValueDict()
+    for op in startup_program.global_block().ops:
+        if op.name() == "builtin.set_parameter":
+            name = op.attrs()["parameter_name"]
+            param = op.operand(0).source()
+            startup_params[name] = param
+
+    for op in main_program.global_block().ops:
+        if op.name() == "builtin.parameter":
+            name = op.attrs()["parameter_name"]
+            param = op.result(0)
+            main_params[name] = param
+
+    assert len(startup_params) == len(main_params)
+    for name, startup_param in startup_params.items():
+        assert name in main_params
+        main_param = main_params[name]
+        parameter_mapping[main_param] = startup_param
+    return parameter_mapping
+
+
 class SimpleImgConvPool(paddle.nn.Layer):
     def __init__(
         self,
@@ -208,16 +233,30 @@ def test_mnist_float32(self):
             # initialize params and fetch them
             static_param_init_value = {}
             static_param_name_list = []
+            static_params = []
             for param in mnist.parameters():
                 static_param_name_list.append(param.name)
+                static_params.append(param)
+
+            if paddle.framework.use_pir_api():
+                parameter_mapping = create_parameter_mapping(
+                    paddle.static.default_startup_program(),
+                    paddle.static.default_main_program(),
+                )
+                startup_params = [
+                    parameter_mapping[param] for param in static_params
+                ]
+            else:
+                startup_params = static_params
 
             out = exe.run(
-                base.default_startup_program(),
-                fetch_list=static_param_name_list,
+                paddle.static.default_startup_program(),
+                fetch_list=startup_params,
             )
 
-            for i in range(len(static_param_name_list)):
-                static_param_init_value[static_param_name_list[i]] = out[i]
+            for i in range(len(static_params)):
+                param_name = static_param_name_list[i]
+                static_param_init_value[param_name] = out[i]
 
             for epoch in range(epoch_num):
                 for batch_id, data in enumerate(train_reader()):
@@ -232,8 +271,8 @@ def test_mnist_float32(self):
                         .reshape([batch_size, 1])
                     )
 
-                    fetch_list = [avg_loss.name]
-                    fetch_list.extend(static_param_name_list)
+                    fetch_list = [avg_loss]
+                    fetch_list.extend(static_params)
 
                     if traced_layer is not None:
                         traced_layer([static_x_data])
diff --git a/test/legacy_test/test_imperative_mnist_sorted_gradient.py b/test/legacy_test/test_imperative_mnist_sorted_gradient.py
@@ -20,9 +20,34 @@
 
 import paddle
 from paddle import base
+from paddle.autograd.backward_utils import ValueDict
 from paddle.base import core
 
 
+def create_parameter_mapping(startup_program, main_program):
+    startup_params = {}
+    main_params = {}
+    parameter_mapping = ValueDict()
+    for op in startup_program.global_block().ops:
+        if op.name() == "builtin.set_parameter":
+            name = op.attrs()["parameter_name"]
+            param = op.operand(0).source()
+            startup_params[name] = param
+
+    for op in main_program.global_block().ops:
+        if op.name() == "builtin.parameter":
+            name = op.attrs()["parameter_name"]
+            param = op.result(0)
+            main_params[name] = param
+
+    assert len(startup_params) == len(main_params)
+    for name, startup_param in startup_params.items():
+        assert name in main_params
+        main_param = main_params[name]
+        parameter_mapping[main_param] = startup_param
+    return parameter_mapping
+
+
 class TestImperativeMnistSortGradient(unittest.TestCase):
     def test_mnist_sort_gradient_float32(self):
         seed = 90
@@ -110,16 +135,31 @@ def test_mnist_sort_gradient_float32(self):
             # initialize params and fetch them
             static_param_init_value = {}
             static_param_name_list = []
+            static_params = []
+
             for param in mnist.parameters():
                 static_param_name_list.append(param.name)
+                static_params.append(param)
+
+            if paddle.framework.use_pir_api():
+                parameter_mapping = create_parameter_mapping(
+                    paddle.static.default_startup_program(),
+                    paddle.static.default_main_program(),
+                )
+                startup_params = [
+                    parameter_mapping[param] for param in static_params
+                ]
+            else:
+                startup_params = static_params
 
             out = exe.run(
-                base.default_startup_program(),
-                fetch_list=static_param_name_list,
+                paddle.static.default_startup_program(),
+                fetch_list=startup_params,
             )
 
             for i in range(len(static_param_name_list)):
-                static_param_init_value[static_param_name_list[i]] = out[i]
+                param_name = static_param_name_list[i]
+                static_param_init_value[param_name] = out[i]
 
             for epoch in range(epoch_num):
                 for batch_id, data in enumerate(train_reader()):
@@ -132,8 +172,8 @@ def test_mnist_sort_gradient_float32(self):
                         .reshape([128, 1])
                     )
 
-                    fetch_list = [avg_loss.name]
-                    fetch_list.extend(static_param_name_list)
+                    fetch_list = [avg_loss]
+                    fetch_list.extend(static_params)
                     out = exe.run(
                         base.default_main_program(),
                         feed={"pixel": static_x_data, "label": y_data},
@@ -167,4 +207,5 @@ def test_mnist_sort_gradient_float32(self):
 
 
 if __name__ == '__main__':
+    paddle.enable_static()
     unittest.main()
diff --git a/test/legacy_test/test_imperative_reinforcement.py b/test/legacy_test/test_imperative_reinforcement.py
diff --git a/test/legacy_test/test_imperative_transformer_sorted_gradient.py b/test/legacy_test/test_imperative_transformer_sorted_gradient.py
diff --git a/test/legacy_test/test_inverse_op.py b/test/legacy_test/test_inverse_op.py
