Skip to content

Commit 7238195

Browse files
megeminimegemini
authored
[Typing][B-74] Add type annotations for python/paddle/sparse/nn/layer/pooling.py (#65872)
1 parent 97ddf70 commit 7238195

File tree

2 files changed

+124
-11
lines changed

2 files changed

+124
-11
lines changed

‎python/paddle/nn/layer/pooling.py‎

Lines changed: 93 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,13 @@ class AvgPool1D(Layer):
9696
9797
"""
9898

99+
kernel_size: Size1
100+
stride: Size1 | None
101+
padding: _PaddingSizeMode | Size1 | Size2
102+
exclusive: bool
103+
ceil_mode: bool
104+
name: str | None
105+
99106
def __init__(
100107
self,
101108
kernel_size: Size1,
@@ -203,6 +210,15 @@ class AvgPool2D(Layer):
203210
204211
"""
205212

213+
kernel_size: Size2
214+
stride: Size2 | None
215+
padding: _PaddingSizeMode | Size2 | Size4
216+
ceil_mode: bool
217+
exclusive: bool
218+
divisor: float | None
219+
data_format: DataLayout2D
220+
name: str | None
221+
206222
def __init__(
207223
self,
208224
kernel_size: Size2,
@@ -302,6 +318,15 @@ class AvgPool3D(Layer):
302318
303319
"""
304320

321+
kernel_size: Size3
322+
stride: Size3 | None
323+
padding: _PaddingSizeMode | Size3 | Size6
324+
ceil_mode: bool
325+
exclusive: bool
326+
divisor: float | None
327+
data_format: DataLayout3D
328+
name: str | None
329+
305330
def __init__(
306331
self,
307332
kernel_size: Size3,
@@ -402,6 +427,14 @@ class LPPool1D(Layer):
402427
403428
"""
404429

430+
norm_type: float
431+
kernel_size: Size1
432+
stride: Size1 | None
433+
padding: _PaddingSizeMode | Size1 | Size2
434+
ceil_mode: bool
435+
data_format: DataLayout1D
436+
name: str | None
437+
405438
def __init__(
406439
self,
407440
norm_type: float,
@@ -509,6 +542,14 @@ class LPPool2D(Layer):
509542
510543
"""
511544

545+
norm_type: float
546+
kernel_size: Size2
547+
stride: Size2 | None
548+
padding: _PaddingSizeMode | Size2 | Size4
549+
ceil_mode: bool
550+
data_format: DataLayout2D
551+
name: str | None
552+
512553
def __init__(
513554
self,
514555
norm_type: float,
@@ -610,6 +651,13 @@ class MaxPool1D(Layer):
610651
611652
"""
612653

654+
kernel_size: Size1
655+
stride: Size1 | None
656+
padding: _PaddingSizeMode | Size1 | Size2
657+
return_mask: bool
658+
ceil_mode: bool
659+
name: str | None
660+
613661
def __init__(
614662
self,
615663
kernel_size: Size1,
@@ -720,6 +768,14 @@ class MaxPool2D(Layer):
720768
[1, 3, 16, 16]
721769
"""
722770

771+
kernel_size: Size2
772+
stride: Size2 | None
773+
padding: _PaddingSizeMode | Size2 | Size4
774+
return_mask: bool
775+
ceil_mode: bool
776+
data_format: DataLayout2D
777+
name: str | None
778+
723779
def __init__(
724780
self,
725781
kernel_size: Size2,
@@ -820,14 +876,22 @@ class MaxPool3D(Layer):
820876
[1, 2, 1, 16, 16]
821877
"""
822878

879+
kernel_size: Size3
880+
stride: Size3 | None
881+
padding: _PaddingSizeMode | Size3 | Size6
882+
return_mask: bool
883+
ceil_mode: bool
884+
data_format: DataLayout3D
885+
name: str | None
886+
823887
def __init__(
824888
self,
825889
kernel_size: Size3,
826890
stride: Size3 | None = None,
827891
padding: _PaddingSizeMode | Size3 | Size6 = 0,
828892
return_mask: bool = False,
829893
ceil_mode: bool = False,
830-
data_format: DataLayout3D | str = 'NCDHW',
894+
data_format: DataLayout3D = 'NCDHW',
831895
name: str | None = None,
832896
) -> None:
833897
super().__init__()
@@ -908,6 +972,9 @@ class AdaptiveAvgPool1D(Layer):
908972
[1, 3, 16]
909973
"""
910974

975+
output_size: int
976+
name: str | None
977+
911978
def __init__(self, output_size: int, name: str | None = None) -> None:
912979
super().__init__()
913980
self.output_size = output_size
@@ -1178,6 +1245,10 @@ class AdaptiveMaxPool1D(Layer):
11781245
11791246
"""
11801247

1248+
output_size: int
1249+
return_mask: bool
1250+
name: str | None
1251+
11811252
def __init__(
11821253
self,
11831254
output_size: int,
@@ -1446,6 +1517,13 @@ class MaxUnPool1D(Layer):
14461517
14471518
"""
14481519

1520+
kernel_size: Size1
1521+
stride: Size1 | None
1522+
padding: _PaddingSizeMode | Size1 | Size2
1523+
data_format: DataLayout1D
1524+
output_size: Sequence[int] | None
1525+
name: str | None
1526+
14491527
def __init__(
14501528
self,
14511529
kernel_size: Size1,
@@ -1541,6 +1619,13 @@ class MaxUnPool2D(Layer):
15411619
15421620
"""
15431621

1622+
kernel_size: Size2
1623+
stride: Size2 | None
1624+
padding: _PaddingSizeMode | Size2 | Size4
1625+
data_format: DataLayout2D
1626+
output_size: Sequence[int] | None
1627+
name: str | None
1628+
15441629
def __init__(
15451630
self,
15461631
kernel_size: Size2,
@@ -1636,6 +1721,13 @@ class MaxUnPool3D(Layer):
16361721
16371722
"""
16381723

1724+
kernel_size: Size3
1725+
stride: Size3 | None
1726+
padding: _PaddingSizeMode | Size3 | Size6
1727+
data_format: DataLayout3D
1728+
output_size: Sequence[int] | None
1729+
name: str | None
1730+
16391731
def __init__(
16401732
self,
16411733
kernel_size: Size3,

‎python/paddle/sparse/nn/layer/pooling.py‎

Lines changed: 31 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -12,10 +12,23 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
from __future__ import annotations
16+
17+
from typing import TYPE_CHECKING
18+
1519
from paddle.nn import Layer
1620

1721
from .. import functional as F
1822

23+
if TYPE_CHECKING:
24+
from paddle import Tensor
25+
from paddle._typing import (
26+
DataLayout3D,
27+
Size3,
28+
Size6,
29+
)
30+
from paddle.nn.functional.common import _PaddingSizeMode
31+
1932

2033
class MaxPool3D(Layer):
2134
"""
@@ -72,16 +85,24 @@ class MaxPool3D(Layer):
7285
[2, 1, 2, 2, 3]
7386
"""
7487

88+
kernel_size: Size3
89+
stride: Size3 | None
90+
padding: _PaddingSizeMode | Size3 | Size6
91+
return_mask: bool
92+
ceil_mode: bool
93+
data_format: DataLayout3D
94+
name: str | None
95+
7596
def __init__(
7697
self,
77-
kernel_size,
78-
stride=None,
79-
padding=0,
80-
return_mask=False,
81-
ceil_mode=False,
82-
data_format="NDHWC",
83-
name=None,
84-
):
98+
kernel_size: Size3,
99+
stride: Size3 | None = None,
100+
padding: _PaddingSizeMode | Size3 | Size6 = 0,
101+
return_mask: bool = False,
102+
ceil_mode: bool = False,
103+
data_format: DataLayout3D = "NDHWC",
104+
name: str | None = None,
105+
) -> None:
85106
super().__init__()
86107
self.ksize = kernel_size
87108
self.stride = stride
@@ -91,7 +112,7 @@ def __init__(
91112
self.data_format = data_format
92113
self.name = name
93114

94-
def forward(self, x):
115+
def forward(self, x: Tensor) -> Tensor:
95116
return F.max_pool3d(
96117
x,
97118
kernel_size=self.ksize,
@@ -102,7 +123,7 @@ def forward(self, x):
102123
name=self.name,
103124
)
104125

105-
def extra_repr(self):
126+
def extra_repr(self) -> str:
106127
return 'kernel_size={ksize}, stride={stride}, padding={padding}'.format(
107128
**self.__dict__
108129
)

0 commit comments

Comments
 (0)