[Typing][A-17] Add type annotations for conv layers (#65183)

liyongchao911 · SigureMo · web-flow · commit 4eb608dbdc92 · 2024-06-29T03:38:56.000+08:00
---------

Co-authored-by: SigureMo &lt;sigure.qaq@gmail.com&gt;
diff --git a/python/paddle/nn/functional/common.py b/python/paddle/nn/functional/common.py
@@ -63,7 +63,7 @@
     ]
     _DropoutMode: TypeAlias = Literal['upscale_in_train', 'downscale_in_infer']
     _PaddingTensorMode: TypeAlias = Literal[
-        "constant", "reflect", "replicate", "circular"
+        "zeros", "constant", "reflect", "replicate", "circular"
     ]
     _PaddingSizeMode: TypeAlias = Literal[  # noqa: PYI047
         'valid', 'same', 'VALID', 'SAME'
diff --git a/python/paddle/nn/functional/conv.py b/python/paddle/nn/functional/conv.py
@@ -26,6 +26,7 @@
     DataLayout1D,
     DataLayout2D,
     DataLayout3D,
+    DataLayoutND,
     Size1,
     Size2,
     Size3,
@@ -127,20 +128,20 @@ def _update_padding_nd(padding, channel_last, num_dims):
 
 
 def _conv_nd(
-    x,
-    weight,
-    bias=None,
-    stride=1,
-    padding=0,
+    x: Tensor,
+    weight: Tensor,
+    bias: Tensor | None = None,
+    stride: int | Sequence[int] = 1,
+    padding: _PaddingSizeMode | int | Sequence[int] | Sequence[Size2] = 0,
     padding_algorithm=None,
-    dilation=1,
-    groups=1,
-    data_format="NCHW",
-    channel_dim=1,
-    op_type="conv2d",
-    use_cudnn=True,
-    name=None,
-):
+    dilation: int | Sequence[int] = 1,
+    groups: int = 1,
+    data_format: DataLayoutND = "NCHW",
+    channel_dim: int = 1,
+    op_type: str = "conv2d",
+    use_cudnn: bool = True,
+    name: str | None = None,
+) -> Tensor:
     # Due to the poor performance of NHWC, we transpose the input to NCHW.
     if in_dynamic_or_pir_mode() and op_type == "conv2d":
         pre_bias = _C_ops.conv2d(
@@ -777,7 +778,7 @@ def conv1d_transpose(
     weight: Tensor,
     bias: Tensor | None = None,
     stride: Size1 = 1,
-    padding: _PaddingSizeMode | Size1 | Size2 = 0,
+    padding: _PaddingSizeMode | Size1 | Size2 | Sequence[Size2] = 0,
     output_padding: Size1 = 0,
     groups: int = 1,
     dilation: Size1 = 1,
@@ -1356,7 +1357,7 @@ def conv3d(
     weight: Tensor,
     bias: Tensor | None = None,
     stride: Size3 = 1,
-    padding=0,
+    padding: _PaddingSizeMode | Size3 | Size6 | Sequence[Size2] = 0,
     dilation: Size3 = 1,
     groups: int = 1,
     data_format: DataLayout3D = "NCDHW",
diff --git a/python/paddle/nn/layer/conv.py b/python/paddle/nn/layer/conv.py

Original file line number	Diff line number	Diff line change
`@@ -63,7 +63,7 @@`
`63`	`63`	`]`
`64`	`64`	`_DropoutMode: TypeAlias = Literal['upscale_in_train', 'downscale_in_infer']`
`65`	`65`	`_PaddingTensorMode: TypeAlias = Literal[`
`66`		`- "constant", "reflect", "replicate", "circular"`
	`66`	`+ "zeros", "constant", "reflect", "replicate", "circular"`
`67`	`67`	`]`
`68`	`68`	`_PaddingSizeMode: TypeAlias = Literal[ # noqa: PYI047`
`69`	`69`	`'valid', 'same', 'VALID', 'SAME'`