[Typing][B-41,B-48][BUAA] Add type annotations for python/paddle/autograd/ir_backward.py,python/paddle/cost_model/cost_model.py

python/paddle/autograd/ir_backward.py

@@ -16,6 +16,7 @@
 
 import logging
 import warnings
+from typing import TYPE_CHECKING
 
 import paddle.pir
 from paddle.autograd.backward_utils import (
@@ -59,6 +60,11 @@
 """
 __all__ = ['grad', 'calc_gradient', 'calc_gradient_helper']
 
+if TYPE_CHECKING:
+    from typing import Sequence
+
+    from paddle.base.libpaddle.pir import Value
+
 
 def append_full_like(float_value, copy_value, value, state, backward_ops):
     with paddle.amp.auto_cast(enable=False):
@@ -956,7 +962,12 @@ def create_backward_prune_set(
     return outputs_set, inputs_set, no_gradvar_set
 
 
-def calc_gradient_helper(outputs, inputs, grad_outputs, no_grad_set):
+def calc_gradient_helper(
+    outputs: Value | Sequence[Value],
+    inputs: Value | Sequence[Value],
+    grad_outputs: Value | Sequence[Value | None] | None = None,
+    no_grad_set: set[Value] | None = None,
+) -> dict[Value, list[Value]]:
     block = outputs[0].get_defining_op().get_parent_block()
     state = State(block)
     if all_stop_gradient_true(block):
@@ -1038,7 +1049,12 @@ def calc_gradient_helper(outputs, inputs, grad_outputs, no_grad_set):
     return input_grad_map
 
 
-def calc_gradient(outputs, inputs, grad_outputs, no_grad_set):
+def calc_gradient(
+    outputs: Value | Sequence[Value],
+    inputs: Value | Sequence[Value],
+    grad_outputs: Value | Sequence[Value | None] | None = None,
+    no_grad_set: set[Value] | None = None,
+) -> list[Value]:
     """
     calculate gradient of input
 
@@ -1084,15 +1100,15 @@ def calc_gradient(outputs, inputs, grad_outputs, no_grad_set):
 
 
 def grad(
-    outputs,
-    inputs,
-    grad_outputs=None,
-    retain_graph=None,
-    create_graph=False,
-    only_inputs=True,
-    allow_unused=False,
-    no_grad_vars=None,
-):
+    outputs: Value | Sequence[Value],
+    inputs: Value | Sequence[Value],
+    grad_outputs: Value | Sequence[Value | None] | None = None,
+    retain_graph: bool | None = None,
+    create_graph: bool | None = False,
+    only_inputs: bool | None = True,
+    allow_unused: bool | None = False,
+    no_grad_vars: Value | Sequence[Value] | None = None,
+) -> list[Value]:
     '''
     .. note::
         **This API is ONLY available in imperative mode.**

python/paddle/cost_model/cost_model.py

@@ -11,22 +11,28 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from __future__ import annotations
 
 import json
 import os
+from typing import TYPE_CHECKING, Sequence
 
 import numpy as np
 
 import paddle
 from paddle import static
 from paddle.base import core
 
+if TYPE_CHECKING:
+    from paddle.base.compiler import CompiledProgram
+    from paddle.base.framework import Program
+
 
 class CostModel:
     def __init__(self):
         pass
 
-    def build_program(self):
+    def build_program(self) -> tuple[Program, Program]:
         paddle.enable_static()
 
         main_program = static.Program()
@@ -47,11 +53,11 @@ def build_program(self):
 
     def profile_measure(
         self,
-        startup_program,
-        main_program,
-        device='gpu',
-        fetch_cost_list=['time'],
-    ):
+        startup_program: Program | CompiledProgram,
+        main_program: Program | CompiledProgram,
+        device: str = 'gpu',
+        fetch_cost_list: Sequence = ['time'],
+    ) -> None:
         place = paddle.set_device('gpu')
         x = np.random.random(size=(10, 1)).astype('float32')
         exe = paddle.static.Executor(place)
@@ -64,7 +70,7 @@ def profile_measure(
         cost_model = core.CostModel()
         cost_data = cost_model.ProfileMeasure(device)
 
-    def static_cost_data(self):
+    def static_cost_data(self) -> dict[str, str | float]:
         static_cost_data_path = os.path.join(
             os.path.dirname(__file__), "static_op_benchmark.json"
         )
@@ -74,7 +80,9 @@ def static_cost_data(self):
         # return all static cost data
         return load_dict
 
-    def get_static_op_time(self, op_name, forward=True, dtype="float32"):
+    def get_static_op_time(
+        self, op_name: str, forward: bool = True, dtype: str = "float32"
+    ) -> dict[str, str | float]:
         # if forward is True, return op forward time, otherwise return op backward time.
         if op_name is None:
             raise ValueError(