[Typing][B-74] Add type annotations for python/paddle/sparse/nn/layer/pooling.py (#65872)

megemini · web-flow · commit 7238195b2c93 · 2024-07-11T01:03:06.000+08:00
diff --git a/python/paddle/nn/layer/pooling.py b/python/paddle/nn/layer/pooling.py
@@ -96,6 +96,13 @@ class AvgPool1D(Layer):
 
     """
 
+    kernel_size: Size1
+    stride: Size1 | None
+    padding: _PaddingSizeMode | Size1 | Size2
+    exclusive: bool
+    ceil_mode: bool
+    name: str | None
+
     def __init__(
         self,
         kernel_size: Size1,
@@ -203,6 +210,15 @@ class AvgPool2D(Layer):
 
     """
 
+    kernel_size: Size2
+    stride: Size2 | None
+    padding: _PaddingSizeMode | Size2 | Size4
+    ceil_mode: bool
+    exclusive: bool
+    divisor: float | None
+    data_format: DataLayout2D
+    name: str | None
+
     def __init__(
         self,
         kernel_size: Size2,
@@ -302,6 +318,15 @@ class AvgPool3D(Layer):
 
     """
 
+    kernel_size: Size3
+    stride: Size3 | None
+    padding: _PaddingSizeMode | Size3 | Size6
+    ceil_mode: bool
+    exclusive: bool
+    divisor: float | None
+    data_format: DataLayout3D
+    name: str | None
+
     def __init__(
         self,
         kernel_size: Size3,
@@ -402,6 +427,14 @@ class LPPool1D(Layer):
 
     """
 
+    norm_type: float
+    kernel_size: Size1
+    stride: Size1 | None
+    padding: _PaddingSizeMode | Size1 | Size2
+    ceil_mode: bool
+    data_format: DataLayout1D
+    name: str | None
+
     def __init__(
         self,
         norm_type: float,
@@ -509,6 +542,14 @@ class LPPool2D(Layer):
 
     """
 
+    norm_type: float
+    kernel_size: Size2
+    stride: Size2 | None
+    padding: _PaddingSizeMode | Size2 | Size4
+    ceil_mode: bool
+    data_format: DataLayout2D
+    name: str | None
+
     def __init__(
         self,
         norm_type: float,
@@ -610,6 +651,13 @@ class MaxPool1D(Layer):
 
     """
 
+    kernel_size: Size1
+    stride: Size1 | None
+    padding: _PaddingSizeMode | Size1 | Size2
+    return_mask: bool
+    ceil_mode: bool
+    name: str | None
+
     def __init__(
         self,
         kernel_size: Size1,
@@ -720,6 +768,14 @@ class MaxPool2D(Layer):
             [1, 3, 16, 16]
     """
 
+    kernel_size: Size2
+    stride: Size2 | None
+    padding: _PaddingSizeMode | Size2 | Size4
+    return_mask: bool
+    ceil_mode: bool
+    data_format: DataLayout2D
+    name: str | None
+
     def __init__(
         self,
         kernel_size: Size2,
@@ -820,14 +876,22 @@ class MaxPool3D(Layer):
             [1, 2, 1, 16, 16]
     """
 
+    kernel_size: Size3
+    stride: Size3 | None
+    padding: _PaddingSizeMode | Size3 | Size6
+    return_mask: bool
+    ceil_mode: bool
+    data_format: DataLayout3D
+    name: str | None
+
     def __init__(
         self,
         kernel_size: Size3,
         stride: Size3 | None = None,
         padding: _PaddingSizeMode | Size3 | Size6 = 0,
         return_mask: bool = False,
         ceil_mode: bool = False,
-        data_format: DataLayout3D | str = 'NCDHW',
+        data_format: DataLayout3D = 'NCDHW',
         name: str | None = None,
     ) -> None:
         super().__init__()
@@ -908,6 +972,9 @@ class AdaptiveAvgPool1D(Layer):
             [1, 3, 16]
     """
 
+    output_size: int
+    name: str | None
+
     def __init__(self, output_size: int, name: str | None = None) -> None:
         super().__init__()
         self.output_size = output_size
@@ -1178,6 +1245,10 @@ class AdaptiveMaxPool1D(Layer):
 
     """
 
+    output_size: int
+    return_mask: bool
+    name: str | None
+
     def __init__(
         self,
         output_size: int,
@@ -1446,6 +1517,13 @@ class MaxUnPool1D(Layer):
 
     """
 
+    kernel_size: Size1
+    stride: Size1 | None
+    padding: _PaddingSizeMode | Size1 | Size2
+    data_format: DataLayout1D
+    output_size: Sequence[int] | None
+    name: str | None
+
     def __init__(
         self,
         kernel_size: Size1,
@@ -1541,6 +1619,13 @@ class MaxUnPool2D(Layer):
 
     """
 
+    kernel_size: Size2
+    stride: Size2 | None
+    padding: _PaddingSizeMode | Size2 | Size4
+    data_format: DataLayout2D
+    output_size: Sequence[int] | None
+    name: str | None
+
     def __init__(
         self,
         kernel_size: Size2,
@@ -1636,6 +1721,13 @@ class MaxUnPool3D(Layer):
 
     """
 
+    kernel_size: Size3
+    stride: Size3 | None
+    padding: _PaddingSizeMode | Size3 | Size6
+    data_format: DataLayout3D
+    output_size: Sequence[int] | None
+    name: str | None
+
     def __init__(
         self,
         kernel_size: Size3,
diff --git a/python/paddle/sparse/nn/layer/pooling.py b/python/paddle/sparse/nn/layer/pooling.py
@@ -12,10 +12,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 from paddle.nn import Layer
 
 from .. import functional as F
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle._typing import (
+        DataLayout3D,
+        Size3,
+        Size6,
+    )
+    from paddle.nn.functional.common import _PaddingSizeMode
+
 
 class MaxPool3D(Layer):
     """
@@ -72,16 +85,24 @@ class MaxPool3D(Layer):
             [2, 1, 2, 2, 3]
     """
 
+    kernel_size: Size3
+    stride: Size3 | None
+    padding: _PaddingSizeMode | Size3 | Size6
+    return_mask: bool
+    ceil_mode: bool
+    data_format: DataLayout3D
+    name: str | None
+
     def __init__(
         self,
-        kernel_size,
-        stride=None,
-        padding=0,
-        return_mask=False,
-        ceil_mode=False,
-        data_format="NDHWC",
-        name=None,
-    ):
+        kernel_size: Size3,
+        stride: Size3 | None = None,
+        padding: _PaddingSizeMode | Size3 | Size6 = 0,
+        return_mask: bool = False,
+        ceil_mode: bool = False,
+        data_format: DataLayout3D = "NDHWC",
+        name: str | None = None,
+    ) -> None:
         super().__init__()
         self.ksize = kernel_size
         self.stride = stride
@@ -91,7 +112,7 @@ def __init__(
         self.data_format = data_format
         self.name = name
 
-    def forward(self, x):
+    def forward(self, x: Tensor) -> Tensor:
         return F.max_pool3d(
             x,
             kernel_size=self.ksize,
@@ -102,7 +123,7 @@ def forward(self, x):
             name=self.name,
         )
 
-    def extra_repr(self):
+    def extra_repr(self) -> str:
         return 'kernel_size={ksize}, stride={stride}, padding={padding}'.format(
             **self.__dict__
         )
