[Typing][C-80, C-86] Add type annotations for python/paddle/incubate/{nn/functional/variable_length_memory_efficient_attention, operators/graph_khop_sampler}.py (#67197)

inaomIIsfarell · web-flow · commit 280fcd239180 · 2024-08-15T17:52:25.000+08:00
diff --git a/python/paddle/incubate/nn/functional/variable_length_memory_efficient_attention.py b/python/paddle/incubate/nn/functional/variable_length_memory_efficient_attention.py
@@ -18,24 +18,29 @@
 #
 # This source code is licensed under the BSD license found in the
 # LICENSE file in the root directory of this source tree.
+from __future__ import annotations
 
 import math
+from typing import TYPE_CHECKING
 
 from paddle import _C_ops
 from paddle.framework import LayerHelper, in_dynamic_or_pir_mode
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+
 
 def variable_length_memory_efficient_attention(
-    query,
-    key,
-    value,
-    seq_lens,
-    kv_seq_lens,
-    mask=None,
-    scale=None,
-    causal=False,
-    pre_cache_length=0,
-):
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    seq_lens: Tensor,
+    kv_seq_lens: Tensor,
+    mask: Tensor | None = None,
+    scale: float | None = None,
+    causal: bool = False,
+    pre_cache_length: int = 0,
+) -> Tensor:
     """
     Cutlass Memory Efficient Variable Attention.
     This method requires SM_ARCH in sm70, sm75, sm80.
@@ -46,12 +51,12 @@ def variable_length_memory_efficient_attention(
         value (Tensor): The Value Tensor. Its shape is [batchsize, num_head, seq_len, head_size].
         seq_lens (Tensor): The sequence lengths of the sequences in the batch, used to index query. Its shape is [batchsize, 1].
         kv_seq_lens (Tensor): The sequence lengths of the sequences in the batch, used to index key and value. Its shape is [batchsize, 1].
-        mask (Tensor): The Mask Tensor. Its shape is [batchsize, 1, query_seq_len, key_seq_len].
-        scale (Float): The attention matrix's scale. Default is sqrt(1.0 / head_size).
+        mask (Tensor|None, optional): The Mask Tensor. Its shape is [batchsize, 1, query_seq_len, key_seq_len].
+        scale (Float|None, optional): The attention matrix's scale. Default is sqrt(1.0 / head_size).
         causal (Bool): Whether causal masking is used or not. Default is False.
         pre_cache_length (Int): The length of the pre-cache. Default is 0.
     Returns:
-        Tensor: the output Tensor.
+        Tensor, the output Tensor.
 
     Examples:
         .. code-block:: python
diff --git a/python/paddle/incubate/operators/graph_khop_sampler.py b/python/paddle/incubate/operators/graph_khop_sampler.py
@@ -11,12 +11,54 @@
 # 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
+
+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
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+
+@overload
+def graph_khop_sampler(
+    row: Tensor,
+    colptr: Tensor,
+    input_nodes: Tensor,
+    sample_sizes: list[int] | tuple[int, ...],
+    sorted_eids: Tensor | None = ...,
+    return_eids: Literal[True] = ...,
+    name: str | None = None,
+) -> tuple[Tensor, Tensor, Tensor, Tensor]: ...
+
+
+@overload
+def graph_khop_sampler(
+    row: Tensor,
+    colptr: Tensor,
+    input_nodes: Tensor,
+    sample_sizes: list[int] | tuple[int, ...],
+    sorted_eids: Tensor | None = ...,
+    return_eids: Literal[False] = ...,
+    name: str | None = None,
+) -> tuple[Tensor, Tensor, Tensor]: ...
+
+
+@overload
+def graph_khop_sampler(
+    row: Tensor,
+    colptr: Tensor,
+    input_nodes: Tensor,
+    sample_sizes: list[int] | tuple[int, ...],
+    sorted_eids: Tensor | None = ...,
+    return_eids: bool = ...,
+    name: str | None = None,
+) -> tuple[Tensor, Tensor, Tensor, Tensor] | tuple[Tensor, Tensor, Tensor]: ...
+
 
 def graph_khop_sampler(
     row,
@@ -51,7 +93,7 @@ def graph_khop_sampler(
         sample_sizes (list|tuple): The number of neighbors and number of layers we want
                                    to sample. The data type should be int, and the shape
                                    should only have one dimension.
-        sorted_eids (Tensor, optional): The sorted edge ids, should not be None when `return_eids`
+        sorted_eids (Tensor|None, optional): The sorted edge ids, should not be None when `return_eids`
                               is True. The shape should be [num_edges, 1], and the data
                               type should be the same with `row`. Default is None.
         return_eids (bool, optional): Whether to return the id of the sample edges. Default is False.
@@ -69,7 +111,7 @@ def graph_khop_sampler(
 
     Examples:
         .. code-block:: python
-
+            >>> # doctest: +SKIP('Number of sample edges dismatch, the sample kernel has error.')
             >>> import paddle
 
             >>> row = [3, 7, 0, 9, 1, 4, 2, 9, 3, 9, 1, 9, 7]
@@ -80,7 +122,7 @@ def graph_khop_sampler(
             >>> colptr = paddle.to_tensor(colptr, dtype="int64")
             >>> nodes = paddle.to_tensor(nodes, dtype="int64")
 
-            >>> edge_src, edge_dst, sample_index, reindex_nodes = paddle.incubate.graph_khop_sampler(row, colptr, nodes, sample_sizes, False)
+            >>> edge_src, edge_dst, sample_index, reindex_nodes = paddle.incubate.graph_khop_sampler(row, colptr, nodes, sample_sizes, False)   # type: ignore[operator]
 
     """
 
