🏷️ typing: add typing for pir

paddle-stubs/__init__.pyi

@@ -9,6 +9,7 @@ from . import linalg as linalg
 from . import metric as metric
 from . import nn as nn
 from . import optimizer as optimizer
+from . import pir as pir
 from . import regularizer as regularizer
 from . import signal as signal
 from . import static as static

paddle-stubs/base/__init__.pyi

@@ -21,6 +21,7 @@ from . import initializer as initializer
 from . import install_check as install_check
 from . import io as io
 from . import layers as layers
+from . import libpaddle as libpaddle
 from . import metrics as metrics
 from . import nets as nets
 from . import optimizer as optimizer

paddle-stubs/pir/__init__.pyi

@@ -0,0 +1,23 @@
+from paddle.base.libpaddle.pir import (
+    Block,
+    Operation,
+    OpOperand,
+    OpResult,
+    PassManager,
+    Program,
+    Type,
+    Value,
+    check_unregistered_ops,
+    fake_op_result,
+    is_fake_op_result,
+    register_paddle_dialect,
+    reset_insertion_point_to_end,
+    reset_insertion_point_to_start,
+    set_global_program,
+    set_insertion_point,
+    translate_to_new_ir,
+    translate_to_new_ir_with_param_map,
+)
+
+from . import core
+from .math_op_patch import monkey_patch_opresult

paddle-stubs/pir/core.pyi

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+from paddle import Tensor
+from paddle._typing import DTypeLike, ShapeLike
+from paddle.base.core import VarDesc
+from paddle.base.libpaddle import DataType
+
+from ..base.wrapped_decorator import signature_safe_contextmanager
+
+vartype_to_datatype: dict[VarDesc.VarType, DataType]
+
+np_type_to_paddle_type: dict[str, DataType]
+
+def convert_np_dtype_to_dtype_(np_dtype: str | np.dtype) -> DataType: ...
+def default_startup_program(): ...
+def default_main_program(): ...
+def switch_main_program(program): ...
+def switch_startup_program(program): ...
+@signature_safe_contextmanager
+def program_guard(main_program, startup_program=None): ...
+
+class ParameterMeta:
+    shape: ShapeLike
+    dtype: DTypeLike
+
+def create_parameter(
+    dtype: DTypeLike,
+    shape: ShapeLike,
+    **kwargs: Any,
+): ...
+def _convert_into_opresult(tensor: Tensor) -> None: ...

paddle-stubs/pir/math_op_patch.py

@@ -0,0 +1,5 @@
+from __future__ import annotations
+
+
+def monkey_patch_opresult():
+    ...