[Typing][A-15] Add type annotations for paddle/nn/layer/common.py (#65197)

megemini · web-flow · commit c106cfeacc8b · 2024-06-27T20:19:16.000+08:00
diff --git a/python/paddle/_typing/__init__.py b/python/paddle/_typing/__init__.py
@@ -24,7 +24,6 @@
     NestedStructure as NestedStructure,
     Numberic as Numberic,
     NumbericSequence as NumbericSequence,
-    PaddingMode as PaddingMode,
     TensorIndex as TensorIndex,
     TensorLike as TensorLike,
     TensorOrTensors as TensorOrTensors,
diff --git a/python/paddle/_typing/basic.py b/python/paddle/_typing/basic.py
@@ -18,7 +18,6 @@
     Any,
     Dict,
     List,
-    Literal,
     Sequence,
     Tuple,
     TypeVar,
@@ -58,5 +57,3 @@
 NumbericSequence = Sequence[Numberic]
 NestedNumbericSequence: TypeAlias = NestedSequence[Numberic]
 TensorOrTensors: TypeAlias = Union["Tensor", Sequence["Tensor"]]
-
-PaddingMode: TypeAlias = Literal['valid', 'same', 'VALID', 'SAME']
diff --git a/python/paddle/nn/functional/common.py b/python/paddle/nn/functional/common.py
@@ -16,23 +16,10 @@
 
 from typing import TYPE_CHECKING, Literal
 
-if TYPE_CHECKING:
-    from paddle import Tensor
-    from paddle.distributed.communication.group import Group
-
 import numpy
 
 import paddle
 from paddle import _C_ops, pir
-from paddle._typing import (
-    DataLayout2D,
-    DataLayout3D,
-    DataLayoutND,
-    IntSequence,
-    ShapeLike,
-    Size2,
-    Size4,
-)
 from paddle.base.layer_helper import LayerHelper
 from paddle.common_ops_import import Variable, default_main_program
 from paddle.framework import (
@@ -55,14 +42,41 @@
 # TODO: define the common functions to build a neural network
 from ...tensor.manipulation import squeeze, unsqueeze
 
+if TYPE_CHECKING:
+    from typing_extensions import TypeAlias
+
+    from paddle import Tensor
+    from paddle._typing import (
+        DataLayout1DVariant,
+        DataLayout2D,
+        DataLayout3D,
+        DataLayoutND,
+        IntSequence,
+        ShapeLike,
+        Size2,
+        Size4,
+    )
+    from paddle.distributed.communication.group import Group
+
+    _InterpolateMode: TypeAlias = Literal[
+        'linear', 'area', 'nearest', 'bilinear', 'bicubic', 'trilinear'
+    ]
+    _DropoutMode: TypeAlias = Literal['upscale_in_train', 'downscale_in_infer']
+    _PaddingTensorMode: TypeAlias = Literal[
+        "constant", "reflect", "replicate", "circular"
+    ]
+    _PaddingSizeMode: TypeAlias = Literal[  # noqa: PYI047
+        'valid', 'same', 'VALID', 'SAME'
+    ]
+
 __all__ = []
 
 
 def unfold(
     x: Tensor,
     kernel_sizes: Size2,
     strides: Size2 = 1,
-    paddings: Size4 = 0,
+    paddings: Size2 | Size4 = 0,
     dilations: Size2 = 1,
     name: str | None = None,
 ) -> Tensor:
@@ -203,14 +217,12 @@ def unfold(
 def interpolate(
     x: Tensor,
     size: ShapeLike | None = None,
-    scale_factor: ShapeLike | None = None,
-    mode: Literal[
-        'linear', 'area', 'nearest', 'bilinear', 'bicubic', 'trilinear'
-    ] = 'nearest',
+    scale_factor: ShapeLike | float | None = None,
+    mode: _InterpolateMode = 'nearest',
     align_corners: bool = False,
     align_mode: int = 0,
     data_format: (
-        Literal["NCW", "NWC", "NCHW", "NHWC", "NCDHW", "NDHWC"] | None
+        DataLayout1DVariant | DataLayout2D | DataLayout3D | None
     ) = None,
     name: str | None = None,
 ) -> Tensor:
@@ -763,12 +775,12 @@ def upsample(
     x: Tensor,
     size: ShapeLike | None = None,
     scale_factor: ShapeLike | None = None,
-    mode: Literal[
-        'linear', 'nearest', 'bilinear', 'bicubic', 'trilinear'
-    ] = 'nearest',
+    mode: _InterpolateMode = 'nearest',
     align_corners: bool = False,
     align_mode: int = 0,
-    data_format: Literal["NCW", "NWC", "NCHW", "NHWC", "NCDHW", "NDHWC"] = None,
+    data_format: (
+        DataLayout1DVariant | DataLayout2D | DataLayout3D | None
+    ) = None,
     name: str | None = None,
 ) -> Tensor:
     """
@@ -789,6 +801,7 @@ def upsample(
     - 'trilinear' : Trilinear interpolation
     - 'nearest' : Nearest neighbor interpolation
     - 'bicubic' : Bicubic interpolation
+    - 'area': Area interpolation
 
     Linear interpolation is the method of using a line connecting two known quantities
     to determine the value of an unknown quantity between the two known quantities.
@@ -921,7 +934,7 @@ def upsample(
              And :attr:`size` has a higher priority than :attr:`scale_factor`.Has to match input size if
              it is either a list or a tuple or a Tensor. If a list/tuple, each element can be an integer or a Tensor of shape: [1] or [].
              Default: None.
-        mode (str, optional): The resample method. It supports 'linear', 'nearest', 'bilinear',
+        mode (str, optional): The resample method. It supports 'linear', 'nearest', 'bilinear', 'area',
                        'bicubic' and 'trilinear' currently. Default: 'nearest'
         align_corners(bool, optional) :  An optional bool, If True, the centers of the 4 corner pixels of the
                                input and output tensors are aligned, preserving the values at the
@@ -1029,9 +1042,7 @@ def dropout(
     p: float = 0.5,
     axis: int | IntSequence | None = None,
     training: bool = True,
-    mode: Literal[
-        'upscale_in_train', 'downscale_in_infer'
-    ] = "upscale_in_train",
+    mode: _DropoutMode = "upscale_in_train",
     name: str | None = None,
 ) -> Tensor:
     r"""
@@ -1610,7 +1621,7 @@ def alpha_dropout(
 def pad(
     x: Tensor,
     pad: ShapeLike,
-    mode: Literal["constant", "reflect", "replicate", "circular"] = 'constant',
+    mode: _PaddingTensorMode = 'constant',
     value: float = 0.0,
     data_format: DataLayoutND = "NCHW",
     name: str | None = None,
diff --git a/python/paddle/nn/functional/conv.py b/python/paddle/nn/functional/conv.py
@@ -16,11 +16,11 @@
 
 from typing import TYPE_CHECKING, Sequence
 
-from ..._typing.basic import PaddingMode
-
 if TYPE_CHECKING:
     from paddle import Tensor
 
+    from .common import _PaddingSizeMode
+
 from paddle import _C_ops, _legacy_C_ops, get_flags, in_dynamic_mode, pir
 from paddle._typing import (
     DataLayout1D,
@@ -291,7 +291,7 @@ def conv1d(
     weight: Tensor,
     bias: Tensor | None = None,
     stride: Size1 = 1,
-    padding: PaddingMode | Size1 | Size2 | Sequence[Size2] = 0,
+    padding: _PaddingSizeMode | Size1 | Size2 | Sequence[Size2] = 0,
     dilation: Size1 = 1,
     groups: int = 1,
     data_format: DataLayout1D = 'NCL',
@@ -350,7 +350,7 @@ def conv1d(
         bias (Tensor, optional): The bias with shape [M,]. Default: None.
         stride (int|list|tuple, optional): The stride size. If stride is a list/tuple, it must
             contain one integers, (stride_size). Default: 1.
-        padding(int|str|tuple|list, optional): The padding size. Padding could be in one of the following forms.
+        padding (int|str|tuple|list, optional): The padding size. Padding could be in one of the following forms.
             1. a string in ['valid', 'same'].
             2. an int, which means the feature map is zero paded by size of `padding` on both sides.
             3. a list[int] or tuple[int] whose length is 1, which means the feature map is zero paded by size of `padding[0]` on both sides.
@@ -536,7 +536,7 @@ def conv2d(
     weight: Tensor,
     bias: Tensor | None = None,
     stride: Size2 = 1,
-    padding: PaddingMode | Size2 | Size4 | Sequence[Size2] = 0,
+    padding: _PaddingSizeMode | Size2 | Size4 | Sequence[Size2] = 0,
     dilation: Size2 = 1,
     groups: int = 1,
     data_format: DataLayout2D = "NCHW",
@@ -777,7 +777,7 @@ def conv1d_transpose(
     weight: Tensor,
     bias: Tensor | None = None,
     stride: Size1 = 1,
-    padding: PaddingMode | Size1 | Size2 = 0,
+    padding: _PaddingSizeMode | Size1 | Size2 = 0,
     output_padding: Size1 = 0,
     groups: int = 1,
     dilation: Size1 = 1,
@@ -1049,7 +1049,7 @@ def conv2d_transpose(
     weight: Tensor,
     bias: Tensor | None = None,
     stride: Size2 = 1,
-    padding: PaddingMode | Size2 | Size4 | Sequence[Size2] = 0,
+    padding: _PaddingSizeMode | Size2 | Size4 | Sequence[Size2] = 0,
     output_padding: Size2 = 0,
     dilation: Size2 = 1,
     groups: int = 1,
@@ -1533,7 +1533,7 @@ def conv3d_transpose(
     weight: Tensor,
     bias: Tensor | None = None,
     stride: Size3 = 1,
-    padding: PaddingMode | Size3 | Size6 | Sequence[Size2] = 0,
+    padding: _PaddingSizeMode | Size3 | Size6 | Sequence[Size2] = 0,
     output_padding: Size3 = 0,
     groups: int = 1,
     dilation: Size3 = 1,
diff --git a/python/paddle/nn/functional/pooling.py b/python/paddle/nn/functional/pooling.py
@@ -39,8 +39,9 @@
 
 if TYPE_CHECKING:
     from paddle import Tensor
+    from paddle._typing import Size1, Size2
 
-    from ..._typing import PaddingMode, Size1, Size2
+    from .common import _PaddingSizeMode
 
 __all__ = []
 
@@ -185,7 +186,7 @@ def avg_pool1d(
     x: Tensor,
     kernel_size: Size1,
     stride: Size1 | None = None,
-    padding: PaddingMode | Size1 | Size2 = 0,
+    padding: _PaddingSizeMode | Size1 | Size2 = 0,
     exclusive: bool = True,
     ceil_mode: bool = False,
     name: str | None = None,
diff --git a/python/paddle/nn/layer/common.py b/python/paddle/nn/layer/common.py
diff --git a/python/paddle/sparse/nn/layer/conv.py b/python/paddle/sparse/nn/layer/conv.py
