[Typing][C-88,C-89] Add type annotations for python/paddle/incubate/operators/* (#67104)

Sekiro-x · web-flow · commit f0cde33eac70 · 2024-08-10T12:35:18.000+08:00
diff --git a/python/paddle/incubate/operators/graph_sample_neighbors.py b/python/paddle/incubate/operators/graph_sample_neighbors.py
@@ -12,12 +12,61 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal, overload
+
 from paddle import _C_ops
 from paddle.base.data_feeder import check_variable_and_dtype
 from paddle.base.layer_helper import LayerHelper
 from paddle.framework import in_dynamic_or_pir_mode
 from paddle.utils import deprecated
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+@overload
+def graph_sample_neighbors(
+    row: Tensor,
+    colptr: Tensor,
+    input_nodes: Tensor,
+    eids: Tensor | None = ...,
+    perm_buffer: Tensor | None = ...,
+    sample_size: int = ...,
+    return_eids: Literal[True] = ...,
+    flag_perm_buffer: bool = ...,
+    name: str | None = ...,
+) -> tuple[Tensor, Tensor, Tensor]: ...
+
+
+@overload
+def graph_sample_neighbors(
+    row: Tensor,
+    colptr: Tensor,
+    input_nodes: Tensor,
+    eids: Tensor | None = ...,
+    perm_buffer: Tensor | None = ...,
+    sample_size: int = ...,
+    return_eids: Literal[False] = ...,
+    flag_perm_buffer: bool = ...,
+    name: str | None = ...,
+) -> tuple[Tensor, Tensor]: ...
+
+
+@overload
+def graph_sample_neighbors(
+    row: Tensor,
+    colptr: Tensor,
+    input_nodes: Tensor,
+    eids: Tensor | None = ...,
+    perm_buffer: Tensor | None = ...,
+    sample_size: int = ...,
+    return_eids: bool = ...,
+    flag_perm_buffer: bool = ...,
+    name: str | None = ...,
+) -> tuple[Tensor, Tensor] | tuple[Tensor, Tensor, Tensor]: ...
+
 
 @deprecated(
     since="2.4.0",
@@ -95,7 +144,7 @@ def graph_sample_neighbors(
             ...     colptr,
             ...     nodes,
             ...     sample_size=sample_size
-            ... )
+            ... ) # type: ignore[operator]
 
     """
 
diff --git a/python/paddle/incubate/operators/graph_send_recv.py b/python/paddle/incubate/operators/graph_send_recv.py
@@ -12,6 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal
+
 import numpy as np
 
 from paddle import _C_ops
@@ -29,6 +33,9 @@
 )
 from paddle.utils import deprecated
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+
 
 @deprecated(
     since="2.4.0",
@@ -37,9 +44,14 @@
     reason="graph_send_recv in paddle.incubate will be removed in future",
 )
 def graph_send_recv(
-    x, src_index, dst_index, pool_type="sum", out_size=None, name=None
-):
-    r"""
+    x: Tensor,
+    src_index: Tensor,
+    dst_index: Tensor,
+    pool_type: Literal["sum", "mean", "max", "min"] = "sum",
+    out_size: int | Tensor | None = None,
+    name: str | None = None,
+) -> Tensor:
+    """
 
     Graph Learning Send_Recv combine operator.
 
@@ -99,7 +111,7 @@ def graph_send_recv(
             >>> indexes = paddle.to_tensor([[0, 1], [1, 2], [2, 1], [0, 0]], dtype="int32")
             >>> src_index = indexes[:, 0]
             >>> dst_index = indexes[:, 1]
-            >>> out = paddle.incubate.graph_send_recv(x, src_index, dst_index, pool_type="sum")
+            >>> out = paddle.incubate.graph_send_recv(x, src_index, dst_index, pool_type="sum") # type: ignore[operator]
             >>> print(out)
             Tensor(shape=[3, 3], dtype=float32, place=Place(cpu), stop_gradient=True,
             [[0. , 2. , 3. ],
@@ -111,7 +123,7 @@ def graph_send_recv(
             >>> src_index = indexes[:, 0]
             >>> dst_index = indexes[:, 1]
             >>> out_size = paddle.max(dst_index) + 1
-            >>> out = paddle.incubate.graph_send_recv(x, src_index, dst_index, pool_type="sum", out_size=out_size)
+            >>> out = paddle.incubate.graph_send_recv(x, src_index, dst_index, pool_type="sum", out_size=out_size) # type: ignore[operator]
             >>> print(out)
             Tensor(shape=[2, 3], dtype=float32, place=Place(cpu), stop_gradient=True,
             [[0. , 2. , 3. ],
@@ -121,7 +133,7 @@ def graph_send_recv(
             >>> indexes = paddle.to_tensor([[0, 1], [2, 1], [0, 0]], dtype="int32")
             >>> src_index = indexes[:, 0]
             >>> dst_index = indexes[:, 1]
-            >>> out = paddle.incubate.graph_send_recv(x, src_index, dst_index, pool_type="sum")
+            >>> out = paddle.incubate.graph_send_recv(x, src_index, dst_index, pool_type="sum") # type: ignore[operator]
             >>> print(out)
             Tensor(shape=[3, 3], dtype=float32, place=Place(cpu), stop_gradient=True,
             [[0. , 2. , 3. ],
