【Type Hints】Paddle 中引入 Tensor stub 文件

python/CMakeLists.txt

@@ -167,8 +167,16 @@ endif()
 add_custom_target(paddle_python ALL
                   DEPENDS ${PADDLE_PYTHON_BUILD_DIR}/.timestamp)
 if(BUILD_WHL_PACKAGE AND NOT WITH_SETUP_INSTALL)
-  add_custom_target(paddle_copy ALL
-                    DEPENDS ${PADDLE_PYTHON_BUILD_DIR}/.timestamp_wheel)
+  add_custom_target(
+    paddle_copy ALL
+    # generate tensor.pyi for type hints
+    COMMAND
+      ${CMAKE_COMMAND} -E env PYTHONPATH=${PADDLE_BINARY_DIR}/python
+      ${PYTHON_EXECUTABLE} ${PADDLE_SOURCE_DIR}/tools/gen_tensor_stub.py
+      --input-file
+      ${PADDLE_SOURCE_DIR}/python/paddle/tensor/tensor.prototype.pyi
+      --output-file ${PADDLE_BINARY_DIR}/python/paddle/tensor/tensor.pyi
+    DEPENDS ${PADDLE_PYTHON_BUILD_DIR}/.timestamp_wheel)
   add_dependencies(paddle_copy paddle_python)
 endif()
 

python/paddle/__init__.py

@@ -11,6 +11,9 @@
 # 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.
+
+import typing
+
 try:
     from paddle.cuda_env import *  # noqa: F403
     from paddle.version import (  # noqa: F401
@@ -69,8 +72,11 @@
     uint8,
 )
 
-Tensor = framework.core.eager.Tensor
-Tensor.__qualname__ = 'Tensor'
+if typing.TYPE_CHECKING:
+    from .tensor.tensor import Tensor
+else:
+    Tensor = framework.core.eager.Tensor
+    Tensor.__qualname__ = 'Tensor'
 
 import paddle.distributed.fleet
 import paddle.text

python/paddle/_typing/__init__.py

@@ -19,6 +19,7 @@
     NestedSequence as NestedSequence,
     Numberic as Numberic,
     NumbericSequence as NumbericSequence,
+    TensorLike as TensorLike,
     TensorOrTensors as TensorOrTensors,
 )
 

python/paddle/_typing/basic.py

@@ -13,15 +13,17 @@
 # limitations under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence, TypeVar, Union
+from typing import TYPE_CHECKING, Any, Sequence, TypeVar, Union
 
 import numpy as np
+import numpy.typing as npt
 from typing_extensions import TypeAlias
 
 if TYPE_CHECKING:
     from paddle import Tensor
 
 Numberic: TypeAlias = Union[int, float, complex, np.number, "Tensor"]
+TensorLike: TypeAlias = Union[npt.NDArray[Any], "Tensor", Numberic]
 
 _T = TypeVar("_T", bound=Numberic)
 _SeqLevel1: TypeAlias = Sequence[_T]

python/paddle/base/core.pyi

@@ -12,23 +12,37 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-class CPUPlace: ...
-
-class CUDAPlace:
+class Place:
+    def custom_device_id(self) -> int: ...
+    def custom_device_type(self) -> str: ...
+    def gpu_device_id(self) -> int: ...
+    def ipu_device_id(self) -> int: ...
+    def is_cpu_place(self) -> bool: ...
+    def is_cuda_pinned_place(self) -> bool: ...
+    def is_custom_place(self) -> bool: ...
+    def is_gpu_place(self) -> bool: ...
+    def is_ipu_place(self) -> bool: ...
+    def is_xpu_place(self) -> bool: ...
+    def set_place(self, place: Place) -> None: ...
+    def xpu_device_id(self) -> int: ...
+
+class CPUPlace(Place): ...
+
+class CUDAPlace(Place):
     def __init__(self, id: int, /) -> None: ...
 
-class CUDAPinnedPlace: ...
+class CUDAPinnedPlace(Place): ...
 
-class NPUPlace:
+class NPUPlace(Place):
     def __init__(self, id: int, /) -> None: ...
 
-class IPUPlace: ...
+class IPUPlace(Place): ...
 
-class CustomPlace:
+class CustomPlace(Place):
     def __init__(self, name: str, id: int, /) -> None: ...
 
-class MLUPlace:
+class MLUPlace(Place):
     def __init__(self, id: int, /) -> None: ...
 
-class XPUPlace:
+class XPUPlace(Place):
     def __init__(self, id: int, /) -> None: ...

python/paddle/tensor/tensor.prototype.pyi

@@ -0,0 +1,263 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+# The `Tensor` template for `tools/gen_tensor_stub.py` generates the stub file `tensor.pyi`.
+# Add docstring, attributes, methods and alias with type annotaions for `Tensor`
+# if not conveniently coding in original place (like c++ source file).
+
+from typing import Any, overload
+
+import numpy.typing as npt
+from typing_extensions import TypeAlias
+
+import paddle
+from paddle import _typing
+
+# avoid same name: Tensor.slice
+_Slice: TypeAlias = slice
+
+class Tensor:
+    # annotation: ${tensor_docstring}
+
+    # annotation: ${tensor_attributes}
+
+    # If method defined below, we should make the method's signature complete,
+    # and ignore the signature extracted from `paddle.Tensor`.
+    # `gen_tensor.stub.py` will NOT overwrite the signature below.
+    # If method has docstring (ignoring the spaces), `gen_tensor.stub.py` also will NOT overwrite it.
+
+    # annotation: ${tensor_methods}
+    @overload
+    def __init__(self) -> None: ...
+    @overload
+    def __init__(
+        self, dtype, dims, name: str, type, persistable: bool
+    ) -> None: ...
+    @overload
+    def __init__(
+        self,
+        value: npt.NDArray[Any],
+        place,
+        persistable: bool,
+        zero_copy: bool,
+        name: str,
+        stop_gradient: bool,
+    ) -> None: ...
+    @overload
+    def __init__(self, value: npt.NDArray[Any]) -> None: ...
+    @overload
+    def __init__(self, value: Tensor) -> None: ...
+    @overload
+    def __init__(
+        self, value: Tensor, place, name: str, process_mesh, placements
+    ) -> None: ...
+    @overload
+    def __init__(
+        self, value: Tensor, dims, name: str, process_mesh, placements
+    ) -> None: ...
+    @overload
+    def __init__(self, value: Tensor, place, name: str) -> None: ...
+    @overload
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        """
+        ref: paddle/fluid/pybind/eager.cc
+
+        We should have init function with signature:
+        1.
+        def __init__ ()
+        2.
+        def __init__ (
+            dtype: paddle::framework::proto::VarType::Type,
+            dims: vector<int>,
+            name: std::string,
+            type: paddle::framework::proto::VarType::LodTensor,
+            persistable: bool)
+        3. (multi-place)
+        (should have at least one parameter, one parameter equals to case 4, zero
+        parameter equals to case 1)
+        def __init__ (
+            value: ndarray,
+            place: paddle::platform::Place,
+            persistable: bool,
+            zero_copy: bool,
+            name: std::string,
+            stop_gradient: bool)
+        4.
+        def __init__ (
+            value: ndarray)
+        5.
+        def __init__ (
+            tensor: Tensor)
+        6. (multi-place)
+        (should have at least one parameter, one parameter equals to case 5, zero
+        parameter equals to case 1.)
+        def __init__ (
+            global_tensor: Tensor,
+            place: paddle::platform::Place,
+            name: std::string,
+            process_mesh: phi::distributed::ProcessMesh
+            placements: std::vector<Placement>)
+        7. (multi-place)
+        (should have at least one parameter, one parameter equals to case 5, zero
+        parameter equals to case 1.)
+        def __init__ (
+            local_tensor: Tensor,
+            global_dims: vector<int>,
+            name: std::string,
+            process_mesh: phi::distributed::ProcessMesh
+            placements: std::vector<Placement>)
+        8. (multi-place) (should have at least one parameter, one parameter similar
+        to case 5, zero parameter equals to case 1.)
+        def __init__ (
+            tensor: FrameworkTensor,
+            place: paddle::platform::Place,
+            name: std::string)
+        """
+        ...
+    # rich comparison
+    def __eq__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore[override]
+    def __ge__(self, y: _typing.TensorLike) -> Tensor: ...
+    def __gt__(self, y: _typing.TensorLike) -> Tensor: ...
+    def __lt__(self, y: _typing.TensorLike) -> Tensor: ...
+    def __le__(self, y: _typing.TensorLike) -> Tensor: ...
+    def __ne__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore[override]
+
+    # binary arithmetic operations
+    def __add__(self, y: _typing.TensorLike) -> Tensor: ...
+    def __sub__(self, y: _typing.TensorLike) -> Tensor: ...
+    def __mul__(self, y: _typing.TensorLike) -> Tensor: ...
+    def __matmul__(self, y: _typing.TensorLike) -> Tensor: ...
+    def __truediv__(self, y: _typing.TensorLike) -> Tensor: ...
+    def __floordiv__(self, y: _typing.TensorLike) -> Tensor: ...
+    def __mod__(self, y: _typing.TensorLike) -> Tensor: ...
+    def __pow__(self, y: _typing.TensorLike) -> Tensor: ...
+    def __and__(self, y: _typing.TensorLike) -> Tensor: ...
+    def __div__(self, y: _typing.TensorLike) -> Tensor: ...
+    def __radd__(self, y: _typing.TensorLike) -> Tensor: ...
+    def __rsub__(self, y: _typing.TensorLike) -> Tensor: ...
+    def __rmul__(self, y: _typing.TensorLike) -> Tensor: ...
+    def __rtruediv__(self, y: _typing.TensorLike) -> Tensor: ...
+    def __rpow__(self, y: _typing.TensorLike) -> Tensor: ...
+    def __rdiv__(self, y: _typing.TensorLike) -> Tensor: ...
+
+    # type cast
+    def __bool__(self) -> bool: ...
+    def __float__(self) -> float: ...
+    def __int__(self) -> int: ...
+    def __long__(self) -> float: ...
+    def __nonzero__(self) -> bool: ...
+
+    # emulating container types
+    def __getitem__(
+        self,
+        item: (
+            None
+            | bool
+            | int
+            | _Slice
+            | tuple[None | bool | int | _Slice, ...]
+            | list[Tensor | bool | int]
+        ),
+    ) -> Tensor: ...
+    def __setitem__(
+        self,
+        item: (
+            None
+            | bool
+            | int
+            | _Slice
+            | tuple[None | bool | int | _Slice, ...]
+            | list[Tensor | bool | int]
+        ),
+        value: Tensor | npt.NDArray[Any] | int | float | complex | bool,
+    ) -> None: ...
+    def __len__(self) -> int: ...
+
+    # emulating numeric types
+    def __index__(self) -> int: ...
+
+    # unary arithmetic operations
+    def __invert__(self) -> Tensor: ...
+    def __neg__(self) -> Tensor: ...
+
+    # basic
+    def __hash__(self) -> int: ...
+    def clear_gradient(self, set_to_zero: bool = True) -> None: ...
+    def clone(self) -> Tensor: ...
+    def cols(self) -> Tensor: ...
+    def contiguous(self) -> Tensor: ...
+    def copy_(self) -> Tensor: ...
+    def crows(self) -> Tensor: ...
+    @property
+    def data(self) -> Tensor: ...
+    def data_ptr(self) -> int: ...
+    def detach(self) -> Tensor: ...
+    def detach_(self) -> Tensor: ...
+    @property
+    def dtype(self) -> paddle.dtype: ...
+    def element_size(self) -> int: ...
+    def get_map_tensor(self) -> Tensor: ...
+    def get_selected_rows(self) -> None: ...
+    def get_strides(self) -> list[int]: ...
+    def get_tensor(self) -> Tensor: ...
+    @property
+    def grad(self) -> Tensor | None: ...
+    @property
+    def grad_(self) -> Tensor | None: ...
+    @property
+    def grad_fn(self) -> Any: ...
+    def is_contiguous(self) -> bool: ...
+    def is_dense(self) -> bool: ...
+    def is_dist(self) -> bool: ...
+    @property
+    def is_leaf(self) -> bool: ...
+    def is_same_shape(self, y: Tensor) -> bool: ...
+    def is_selected_rows(self) -> bool: ...
+    def is_sparse(self) -> bool: ...
+    def is_sparse_coo(self) -> bool: ...
+    def is_sparse_csr(self) -> bool: ...
+    @property
+    def layout(self) -> _typing.DataLayoutND: ...
+    @property
+    def name(self) -> str: ...
+    @property
+    def ndim(self) -> int: ...
+    def nnz(self) -> int: ...
+    @property
+    def num_shard(self) -> int: ...
+    def numpy(self) -> npt.NDArray[Any]: ...
+    @property
+    def offset(self) -> int: ...
+    @property
+    def persistable(self) -> bool: ...
+    @property
+    def place(self) -> paddle.core.Place: ...
+    @property
+    def placements(self) -> list[paddle.distributed.Placement] | None: ...
+    @property
+    def process_mesh(self) -> paddle.distributed.ProcessMesh | None: ...
+    def rows(self) -> list[int]: ...
+    def set_string_list(self, value: str) -> None: ...
+    def set_vocab(self, value: dict) -> None: ...
+    @property
+    def shape(self) -> list[int]: ...
+    @property
+    def size(self) -> int: ...
+    @property
+    def strides(self) -> list[int]: ...
+    @property
+    def type(self) -> Any: ...
+
+    # annotation: ${tensor_alias}
+    __qualname__ = "Tensor"

python/setup.py.in

@@ -1087,6 +1087,8 @@ package_data['paddle.framework'] = package_data.get(
 package_data['paddle.base'] = package_data.get('paddle.base', []) + [
     '*.pyi'
 ]
+package_data['paddle.tensor'] = package_data.get('paddle.tensor', []) + ['tensor.pyi']
+
 
 with redirect_stdout():
     setup(name='${PACKAGE_NAME}',

setup.py

@@ -1318,6 +1318,9 @@ def get_package_data_and_package_dir():
     package_data['paddle.base'] = package_data.get('paddle.base', []) + [
         '*.pyi'
     ]
+    package_data['paddle.tensor'] = package_data.get('paddle.tensor', []) + [
+        'tensor.pyi'
+    ]
 
     return package_data, package_dir, ext_modules