[Typing] 修改 core.pyi 及一些示例代码中的类型错误 (#65496)

megemini · web-flow · commit 1575a1f99310 · 2024-06-28T10:02:42.000+08:00
diff --git a/paddle/fluid/pybind/eager_properties.cc b/paddle/fluid/pybind/eager_properties.cc
@@ -411,12 +411,13 @@ Get process_mesh property from shard tensor.
         >>> # doctest: +REQUIRES(env:DISTRIBUTED)
         >>> import paddle
         >>> import paddle.distributed as dist
+        >>> from paddle.base import core
 
         >>> mesh = dist.ProcessMesh([[2, 4, 5], [0, 1, 3]], dim_names=["x", "y"])
 
         >>> a = paddle.to_tensor([[1,2,3],
         ...                       [5,6,7]])
-        >>> d_tensor = dist.shard_tensor(a, [core.Shard(0), core.Shard(1)])
+        >>> d_tensor = dist.shard_tensor(a, mesh, [core.Shard(0), core.Shard(1)])
 
         >>> print(d_tensor.process_mesh)
 
@@ -463,7 +464,7 @@ Get placements property from shard tensor.
 
         >>> a = paddle.to_tensor([[1,2,3],
         ...                       [5,6,7]])
-        >>> d_tensor = dist.shard_tensor(a, [core.Shard(0), core.Shard(1)])
+        >>> d_tensor = dist.shard_tensor(a, mesh, [core.Shard(0), core.Shard(1)])
 
         >>> print(d_tensor.placements)
 
diff --git a/python/paddle/base/core.pyi b/python/paddle/base/core.pyi
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Any
+
 class Place:
     def custom_device_id(self) -> int: ...
     def custom_device_type(self) -> str: ...
@@ -26,23 +28,55 @@ class Place:
     def set_place(self, place: Place) -> None: ...
     def xpu_device_id(self) -> int: ...
 
-class CPUPlace(Place): ...
+class CPUPlace:
+    def is_float16_supported(self) -> bool: ...
+    def is_bfloat16_supported(self) -> bool: ...
 
-class CUDAPlace(Place):
+class CUDAPlace:
     def __init__(self, id: int, /) -> None: ...
+    def get_device_id(self) -> int: ...
+    def is_float16_supported(self) -> bool: ...
+    def is_bfloat16_supported(self) -> bool: ...
 
-class CUDAPinnedPlace(Place): ...
+class CUDAPinnedPlace: ...
 
-class NPUPlace(Place):
+class NPUPlace:
     def __init__(self, id: int, /) -> None: ...
 
-class IPUPlace(Place): ...
+class IPUPlace: ...
 
-class CustomPlace(Place):
+class CustomPlace:
     def __init__(self, name: str, id: int, /) -> None: ...
+    def get_device_id(self) -> int: ...
+    def get_device_type(self) -> str: ...
 
-class MLUPlace(Place):
+class MLUPlace:
     def __init__(self, id: int, /) -> None: ...
 
-class XPUPlace(Place):
+class XPUPlace:
     def __init__(self, id: int, /) -> None: ...
+    def get_xpu_device_count(self) -> int: ...
+    def set_xpu_debug_level(self, level: int) -> None: ...
+    def get_xpu_device_version(self, device_id: int) -> int: ...
+    def get_xpu_device_op_support_types(
+        self, op_name: str, version: int
+    ) -> list[str]: ...
+    def get_xpu_device_op_list(self, version: int) -> dict[str, Any]: ...
+    def is_float16_supported(self) -> bool: ...
+    def is_bfloat16_supported(self) -> bool: ...
+
+class Placement:
+    def __init__(self) -> None: ...
+    def __eq__(self, arg0: Placement) -> bool: ...
+    def __hash__(self) -> int: ...
+    def __ne__(self, arg0: Placement) -> bool: ...
+    def is_partial(self) -> bool: ...
+    def is_replicated(self) -> bool: ...
+    def is_shard(self, dim: int | None = None) -> bool: ...
+
+class Shard(Placement):
+    def __init__(self, arg0: int) -> None: ...
+    def get_dim(self) -> int: ...
+
+def is_compiled_with_cuda() -> bool: ...
+def set_nan_inf_debug_path(arg0: str) -> None: ...
diff --git a/python/paddle/base/framework.py b/python/paddle/base/framework.py
@@ -8054,7 +8054,7 @@ def device_guard(device=None):
             >>> support_gpu = paddle.is_compiled_with_cuda()
             >>> place = paddle.CPUPlace()
             >>> if support_gpu:
-            ...     place = paddle.CUDAPlace(0)
+            ...     place = paddle.CUDAPlace(0) # type: ignore
 
             >>> # if GPU is supported, the three OPs below will be automatically assigned to CUDAPlace(0)
             >>> data1 = paddle.full(shape=[1, 3, 8, 8], fill_value=0.5, dtype='float32')
diff --git a/python/paddle/distributed/fleet/base/distributed_strategy.py b/python/paddle/distributed/fleet/base/distributed_strategy.py
@@ -1110,6 +1110,14 @@ def recompute(self):
         """
         return self.strategy.recompute
 
+    @recompute.setter
+    @is_strict_auto
+    def recompute(self, flag):
+        if isinstance(flag, bool):
+            self.strategy.recompute = flag
+        else:
+            logger.warning("recompute should have value of bool type")
+
     @property
     def sync_nccl_allreduce(self):
         """
@@ -1367,14 +1375,6 @@ def nccl_comm_num(self, value):
         else:
             logger.warning("nccl_comm_num should have value of int type")
 
-    @recompute.setter
-    @is_strict_auto
-    def recompute(self, flag):
-        if isinstance(flag, bool):
-            self.strategy.recompute = flag
-        else:
-            logger.warning("recompute should have value of bool type")
-
     @property
     def recompute_configs(self):
         """
diff --git a/python/paddle/distributed/fleet/utils/fs.py b/python/paddle/distributed/fleet/utils/fs.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# type: ignore[abstract]
+
 import abc
 import functools
 import multiprocessing
diff --git a/python/paddle/distributed/sharding/group_sharded.py b/python/paddle/distributed/sharding/group_sharded.py
@@ -77,6 +77,7 @@ def group_sharded_parallel(
     Examples:
         .. code-block:: python
 
+            >>> # type: ignore
             >>> # doctest: +REQUIRES(env:DISTRIBUTED)
             >>> import paddle
             >>> from paddle.nn import Linear
@@ -200,6 +201,7 @@ def save_group_sharded_model(model, output, optimizer=None):
     Examples:
         .. code-block:: python
 
+            >>> # type: ignore
             >>> # doctest: +REQUIRES(env:DISTRIBUTED)
             >>> import paddle
             >>> from paddle.nn import Linear
diff --git a/python/paddle/vision/models/mobilenetv3.py b/python/paddle/vision/models/mobilenetv3.py
@@ -94,14 +94,14 @@ def forward(self, input: Tensor) -> Tensor:
 class InvertedResidualConfig:
     def __init__(
         self,
-        in_channels,
-        kernel,
-        expanded_channels,
-        out_channels,
-        use_se,
-        activation,
-        stride,
-        scale=1.0,
+        in_channels: int,
+        kernel: int,
+        expanded_channels: int,
+        out_channels: int,
+        use_se: bool,
+        activation: str,
+        stride: int,
+        scale: float = 1.0,
     ):
         self.in_channels = self.adjust_channels(in_channels, scale=scale)
         self.kernel = kernel
@@ -136,8 +136,8 @@ def __init__(
         filter_size: int,
         stride: int,
         use_se: bool,
-        activation_layer: nn.Layer,
-        norm_layer: nn.Layer,
+        activation_layer: type[nn.Layer],
+        norm_layer: type[nn.Layer],
     ) -> None:
         super().__init__()
         self.use_res_connect = stride == 1 and in_channels == out_channels
diff --git a/python/paddle/vision/models/vgg.py b/python/paddle/vision/models/vgg.py
@@ -118,7 +118,9 @@ def forward(self, x: Tensor) -> Tensor:
         return x
 
 
-def make_layers(cfg: list[object], batch_norm: bool = False) -> Sequential:
+def make_layers(
+    cfg: list[int | Literal['M']], batch_norm: bool = False
+) -> Sequential:
     layers = []
     in_channels = 3
     for v in cfg:
