[Typing][B-72] Add type annotations for python/paddle/sparse/nn/layer/conv.py (#65870)

megemini · web-flow · commit 95b44dcf08d0 · 2024-07-10T09:47:16.000+08:00
diff --git a/python/paddle/sparse/nn/layer/conv.py b/python/paddle/sparse/nn/layer/conv.py
@@ -14,7 +14,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Literal
+from typing import TYPE_CHECKING, Literal, Sequence
 
 import numpy as np
 
@@ -26,33 +26,40 @@
 from .. import functional as F
 
 if TYPE_CHECKING:
-    from typing_extensions import TypeAlias
-
-    _ConvPaddingMode: TypeAlias = Literal[
-        "zero", "reflect", "replicate", "circular"
-    ]
+    from paddle import Tensor
+    from paddle._typing import (
+        ParamAttrLike,
+        Size2,
+        Size3,
+        Size4,
+        Size6,
+    )
+    from paddle.nn.functional.common import _PaddingSizeMode
 
 __all__ = []
 
 
 class _Conv3D(Layer):
+    weight: Tensor
+    bias: Tensor
+
     def __init__(
         self,
-        in_channels,
-        out_channels,
-        kernel_size,
-        stride=1,
-        padding=0,
-        dilation=1,
-        groups=1,
-        subm=False,
-        key=None,
-        padding_mode: _ConvPaddingMode = 'zeros',
-        weight_attr=None,
-        bias_attr=None,
-        data_format="NDHWC",
-        backend=None,
-    ):
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Size3,
+        stride: Size3 = 1,
+        padding: _PaddingSizeMode | Size3 | Size6 | Sequence[Size2] = 0,
+        dilation: Size3 = 1,
+        groups: Literal[1] = 1,
+        subm: bool = False,
+        key: str | None = None,
+        padding_mode: Literal['zeros'] = 'zeros',
+        weight_attr: ParamAttrLike | None = None,
+        bias_attr: ParamAttrLike | None = None,
+        data_format: Literal['NDHWC'] = "NDHWC",
+        backend: Literal['igemm'] | None = None,
+    ) -> None:
         super().__init__()
         assert (
             weight_attr is not False
@@ -114,7 +121,7 @@ def _get_default_param_initializer():
             attr=self._bias_attr, shape=[self._out_channels], is_bias=True
         )
 
-    def forward(self, x):
+    def forward(self, x: Tensor) -> Tensor:
         if self._backend is None:
             out = F.conv._conv3d(
                 x,
@@ -147,7 +154,7 @@ def forward(self, x):
             )
         return out
 
-    def extra_repr(self):
+    def extra_repr(self) -> str:
         main_str = '{_in_channels}, {_out_channels}, kernel_size={_kernel_size}'
         if self._stride != [1] * len(self._stride):
             main_str += ', stride={_stride}'
@@ -164,23 +171,26 @@ def extra_repr(self):
 
 
 class _Conv2D(Layer):
+    weight: Tensor
+    bias: Tensor
+
     def __init__(
         self,
-        in_channels,
-        out_channels,
-        kernel_size,
-        stride=1,
-        padding=0,
-        dilation=1,
-        groups=1,
-        subm=False,
-        key=None,
-        padding_mode='zeros',
-        weight_attr=None,
-        bias_attr=None,
-        data_format="NHWC",
-        backend=None,
-    ):
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Size2,
+        stride: Size2 = 1,
+        padding: _PaddingSizeMode | Size2 | Size4 | Sequence[Size2] = 0,
+        dilation: Size2 = 1,
+        groups: Literal[1] = 1,
+        subm: bool = False,
+        key: str | None = None,
+        padding_mode: Literal['zeros'] = 'zeros',
+        weight_attr: ParamAttrLike | None = None,
+        bias_attr: ParamAttrLike | None = None,
+        data_format: Literal["NHWC"] = "NHWC",
+        backend: Literal['igemm'] | None = None,
+    ) -> None:
         super().__init__()
         assert (
             weight_attr is not False
@@ -242,7 +252,7 @@ def _get_default_param_initializer():
             attr=self._bias_attr, shape=[self._out_channels], is_bias=True
         )
 
-    def forward(self, x):
+    def forward(self, x: Tensor) -> Tensor:
         if self._backend is None:
             out = F.conv._conv2d(
                 x,
@@ -275,7 +285,7 @@ def forward(self, x):
             )
         return out
 
-    def extra_repr(self):
+    def extra_repr(self) -> str:
         main_str = '{_in_channels}, {_out_channels}, kernel_size={_kernel_size}'
         if self._stride != [1] * len(self._stride):
             main_str += ', stride={_stride}'
@@ -397,18 +407,18 @@ class Conv3D(_Conv3D):
 
     def __init__(
         self,
-        in_channels,
-        out_channels,
-        kernel_size,
-        stride=1,
-        padding=0,
-        dilation=1,
-        groups=1,
-        padding_mode='zeros',
-        weight_attr=None,
-        bias_attr=None,
-        data_format="NDHWC",
-    ):
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Size3,
+        stride: Size3 = 1,
+        padding: _PaddingSizeMode | Size3 | Size6 | Sequence[Size2] = 0,
+        dilation: Size3 = 1,
+        groups: Literal[1] = 1,
+        padding_mode: Literal['zeros'] = 'zeros',
+        weight_attr: ParamAttrLike | None = None,
+        bias_attr: ParamAttrLike | None = None,
+        data_format: Literal["NDHWC"] = "NDHWC",
+    ) -> None:
         super().__init__(
             in_channels,
             out_channels,
@@ -532,18 +542,18 @@ class Conv2D(_Conv2D):
 
     def __init__(
         self,
-        in_channels,
-        out_channels,
-        kernel_size,
-        stride=1,
-        padding=0,
-        dilation=1,
-        groups=1,
-        padding_mode='zeros',
-        weight_attr=None,
-        bias_attr=None,
-        data_format="NHWC",
-    ):
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Size2,
+        stride: Size2 = 1,
+        padding: _PaddingSizeMode | Size2 | Size4 | Sequence[Size2] = 0,
+        dilation: Size2 = 1,
+        groups: Literal[1] = 1,
+        padding_mode: Literal['zeros'] = 'zeros',
+        weight_attr: ParamAttrLike | None = None,
+        bias_attr: ParamAttrLike | None = None,
+        data_format: Literal["NHWC"] = "NHWC",
+    ) -> None:
         super().__init__(
             in_channels,
             out_channels,
@@ -671,20 +681,20 @@ class SubmConv3D(_Conv3D):
 
     def __init__(
         self,
-        in_channels,
-        out_channels,
-        kernel_size,
-        stride=1,
-        padding=0,
-        dilation=1,
-        groups=1,
-        padding_mode='zeros',
-        key=None,
-        weight_attr=None,
-        bias_attr=None,
-        data_format="NDHWC",
-        backend=None,
-    ):
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Size3,
+        stride: Size3 = 1,
+        padding: _PaddingSizeMode | Size3 | Size6 | Sequence[Size2] = 0,
+        dilation: Size3 = 1,
+        groups: Literal[1] = 1,
+        padding_mode: Literal['zeros'] = 'zeros',
+        key: str | None = None,
+        weight_attr: ParamAttrLike | None = None,
+        bias_attr: ParamAttrLike | None = None,
+        data_format: Literal["NDHWC"] = "NDHWC",
+        backend: Literal['igemm'] | None = None,
+    ) -> None:
         super().__init__(
             in_channels,
             out_channels,
@@ -813,20 +823,20 @@ class SubmConv2D(_Conv2D):
 
     def __init__(
         self,
-        in_channels,
-        out_channels,
-        kernel_size,
-        stride=1,
-        padding=0,
-        dilation=1,
-        groups=1,
-        padding_mode='zeros',
-        key=None,
-        weight_attr=None,
-        bias_attr=None,
-        data_format="NHWC",
-        backend=None,
-    ):
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Size2,
+        stride: Size2 = 1,
+        padding: _PaddingSizeMode | Size2 | Size4 | Sequence[Size2] = 0,
+        dilation: Size2 = 1,
+        groups: Literal[1] = 1,
+        padding_mode: Literal['zeros'] = 'zeros',
+        key: str | None = None,
+        weight_attr: ParamAttrLike | None = None,
+        bias_attr: ParamAttrLike | None = None,
+        data_format: Literal["NHWC"] = "NHWC",
+        backend: Literal['igemm'] | None = None,
+    ) -> None:
         super().__init__(
             in_channels,
             out_channels,
