[Typing][A-21,A-62] Add type annotations for paddle/nn/{layer,functional}/norm.py (#65454)

Asthestarsfalll · SigureMo · web-flow · commit 365fbd13460e · 2024-06-29T11:36:39.000+08:00
---------

Co-authored-by: Nyakku Shigure &lt;sigure.qaq@gmail.com&gt;
diff --git a/python/paddle/nn/functional/activation.py b/python/paddle/nn/functional/activation.py
@@ -16,12 +16,8 @@
 
 from typing import TYPE_CHECKING
 
-if TYPE_CHECKING:
-    from paddle import Tensor
-
 import paddle
 from paddle import _C_ops, _legacy_C_ops, in_dynamic_mode
-from paddle._typing import DataLayout2D, DTypeLike
 from paddle.framework import core, in_dynamic_or_pir_mode
 from paddle.utils.inplace_utils import inplace_apis_in_dygraph_only
 
@@ -32,6 +28,10 @@
 from ...tensor.math import tanh, tanh_  # noqa: F401
 from ...tensor.ops import sigmoid
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle._typing import DataLayout2D, DTypeLike
+
 __all__ = []
 
 
diff --git a/python/paddle/nn/functional/conv.py b/python/paddle/nn/functional/conv.py
@@ -16,23 +16,7 @@
 
 from typing import TYPE_CHECKING, Sequence
 
-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,
-    DataLayout2D,
-    DataLayout3D,
-    DataLayoutND,
-    Size1,
-    Size2,
-    Size3,
-    Size4,
-    Size6,
-)
 from paddle.base.framework import _global_flags, in_dynamic_or_pir_mode
 from paddle.device import (
     get_all_custom_device_type,
@@ -55,6 +39,23 @@
     convert_to_list,
 )
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle._typing import (
+        DataLayout1D,
+        DataLayout2D,
+        DataLayout3D,
+        DataLayoutND,
+        Size1,
+        Size2,
+        Size3,
+        Size4,
+        Size6,
+    )
+
+    from .common import _PaddingSizeMode
+
+
 __all__ = []
 
 
diff --git a/python/paddle/nn/functional/norm.py b/python/paddle/nn/functional/norm.py
@@ -12,7 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import numbers
+from typing import TYPE_CHECKING, Literal, Sequence
 
 import paddle
 from paddle import _C_ops, in_dynamic_mode
@@ -25,10 +28,25 @@
 from ...base.data_feeder import check_type, check_variable_and_dtype
 from ...base.layer_helper import LayerHelper
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle._typing import (
+        DataLayout1D,
+        DataLayout2D,
+        DataLayout3D,
+        DataLayoutND,
+    )
+
 __all__ = []
 
 
-def normalize(x, p=2, axis=1, epsilon=1e-12, name=None):
+def normalize(
+    x: Tensor,
+    p: float = 2,
+    axis: int = 1,
+    epsilon: float = 1e-12,
+    name: str | None = None,
+) -> Tensor:
     r"""
     Normalize ``x`` along dimension ``axis`` using :math:`L_p` norm. This layer computes
 
@@ -47,7 +65,7 @@ def normalize(x, p=2, axis=1, epsilon=1e-12, name=None):
         p (float|int, optional): The exponent value in the norm formulation. Default: 2.
         axis (int, optional): The axis on which to apply normalization. If `axis < 0`, the dimension to normalization is `x.ndim + axis`. -1 is the last dimension.
         epsilon (float, optional): Small float added to denominator to avoid dividing by zero. Default is 1e-12.
-        name (str, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
+        name (str|None, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
 
     Returns:
         Tensor, the output has the same shape and data type with ``x``.
@@ -120,17 +138,17 @@ def normalize(x, p=2, axis=1, epsilon=1e-12, name=None):
 
 def batch_norm(
     x,
-    running_mean,
-    running_var,
-    weight=None,
-    bias=None,
-    training=False,
-    momentum=0.9,
-    epsilon=1e-05,
-    data_format="NCHW",
-    use_global_stats=None,
-    name=None,
-):
+    running_mean: Tensor,
+    running_var: Tensor,
+    weight: Tensor | None = None,
+    bias: Tensor | None = None,
+    training: bool = False,
+    momentum: float = 0.9,
+    epsilon: float = 1e-05,
+    data_format: DataLayoutND = 'NCHW',
+    use_global_stats: bool | None = None,
+    name: str | None = None,
+) -> Tensor:
     """
     Applies Batch Normalization as described in the paper Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift .
 
@@ -147,7 +165,7 @@ def batch_norm(
         momentum(float, optional): The value used for the moving_mean and moving_var computation. Default: 0.9.
         data_format(str, optional): Specify the input data format, may be "NC", "NCL", "NCHW", "NCDHW", "NLC", "NHWC" or "NDHWC", where `N` is batch size, `C` is the number of the feature map, `D` is the depth of the feature, `H` is the height of the feature map, `W` is the width of the feature map, `L` is the length of the feature map. Default "NCHW".
         use_global_stats(bool|None, optional): Whether to use global mean and variance. If set to False, use the statistics of one mini-batch, if set to True, use the global statistics, if set to None, use global statistics in the test phase and use the statistics of one mini-batch in the training phase. Default: None.
-        name(str, optional): Name for the BatchNorm, default is None. For more information, please refer to :ref:`api_guide_Name`..
+        name(str|None, optional): Name for the BatchNorm, default is None. For more information, please refer to :ref:`api_guide_Name`..
 
     Returns:
         None
@@ -297,8 +315,13 @@ def batch_norm(
 
 
 def layer_norm(
-    x, normalized_shape, weight=None, bias=None, epsilon=1e-05, name=None
-):
+    x: Tensor,
+    normalized_shape: int | Sequence[int],
+    weight: Tensor | None = None,
+    bias: Tensor | None = None,
+    epsilon: float = 1e-05,
+    name: str | None = None,
+) -> Tensor:
     """
     nn.LayerNorm is recommended.
     For more information, please refer to :ref:`api_paddle_nn_LayerNorm` .
@@ -413,33 +436,33 @@ def layer_norm(
 
 
 def instance_norm(
-    x,
-    running_mean=None,
-    running_var=None,
-    weight=None,
-    bias=None,
-    use_input_stats=True,
-    momentum=0.9,
-    eps=1e-05,
-    data_format="NCHW",
-    name=None,
-):
+    x: Tensor,
+    running_mean: Tensor | None = None,
+    running_var: Tensor | None = None,
+    weight: Tensor | None = None,
+    bias: Tensor | None = None,
+    use_input_stats: bool = True,
+    momentum: float = 0.9,
+    eps: float = 1e-05,
+    data_format: Literal['NC', 'NCL', 'NCHW', 'NCDHW'] = 'NCHW',
+    name: str | None = None,
+) -> Tensor:
     """
     It is recommended to use :ref:`api_paddle_nn_InstanceNorm1D` , :ref:`api_paddle_nn_InstanceNorm2D` , :ref:`api_paddle_nn_InstanceNorm3D` to call this method internally.
 
     Parameters:
-        x(Tensor): Input Tensor. It's data type should be float32, float64.
-        running_mean(Tensor, optional): running mean. Default None. Obsolete (that is, no longer usable).
-        running_var(Tensor, optional): running variance. Default None. Obsolete (that is, no longer usable).
-        weight(Tensor, optional): The weight tensor of instance_norm. Default: None.
+        x (Tensor): Input Tensor. It's data type should be float32, float64.
+        running_mean (Tensor, optional): running mean. Default None. Obsolete (that is, no longer usable).
+        running_var (Tensor, optional): running variance. Default None. Obsolete (that is, no longer usable).
+        weight (Tensor, optional): The weight tensor of instance_norm. Default: None.
             If its value is None, this parameter will be initialized by one.
-        bias(Tensor, optional): The bias tensor of instance_norm. Default: None.
+        bias (Tensor, optional): The bias tensor of instance_norm. Default: None.
             If its value is None, this parameter will be initialized by zero.
-        eps(float, optional): A value added to the denominator for numerical stability. Default is 1e-5.
-        momentum(float, optional): The value used for the moving_mean and moving_var computation. Default: 0.9.
-        use_input_stats(bool, optional): Default True. Obsolete (that is, no longer usable).
-        data_format(str, optional): Specify the input data format, may be "NC", "NCL", "NCHW" or "NCDHW". Default "NCHW".
-        name(str, optional): Name for the InstanceNorm, default is None. For more information, please refer to :ref:`api_guide_Name`..
+        eps (float, optional): A value added to the denominator for numerical stability. Default is 1e-5.
+        momentum (float, optional): The value used for the moving_mean and moving_var computation. Default: 0.9.
+        use_input_stats (bool, optional): Default True. Obsolete (that is, no longer usable).
+        data_format (str, optional): Specify the input data format, may be "NC", "NCL", "NCHW" or "NCDHW". Default "NCHW".
+        name (str|None, optional): Name for the InstanceNorm, default is None. For more information, please refer to :ref:`api_guide_Name`..
 
     Returns:
         None.
@@ -509,8 +532,14 @@ def instance_norm(
 
 
 def local_response_norm(
-    x, size, alpha=1e-4, beta=0.75, k=1.0, data_format="NCHW", name=None
-):
+    x: Tensor,
+    size: int,
+    alpha: float = 0.0001,
+    beta: float = 0.75,
+    k: float = 1.0,
+    data_format: DataLayout1D | DataLayout2D | DataLayout3D = 'NCHW',
+    name: str | None = None,
+) -> Tensor:
     r"""
     Local Response Normalization performs a type of "lateral inhibition" by normalizing over local input regions.
     For more information, please refer to `ImageNet Classification with Deep Convolutional Neural Networks <https://papers.nips.cc/paper/4824-imagenet-classification-with-deep-convolutional-neural-networks.pdf>`_
@@ -543,7 +572,7 @@ def local_response_norm(
             the data is stored in the order of: `[batch_size, input_channels, input_height, input_width]`.
             If x is 5-D Tensor, the string could be  `"NCDHW"`, `"NDHWC"` . When it is `"NCDHW"`,
             the data is stored in the order of: `[batch_size, input_channels, input_depth, input_height, input_width]`.
-        name (str, optional): Name for the operation (optional, default is None). For more information,
+        name (str|None, optional): Name for the operation (optional, default is None). For more information,
             please refer to :ref:`api_guide_Name`.
 
     Returns:
@@ -641,14 +670,14 @@ def local_response_norm(
 
 
 def group_norm(
-    x,
-    num_groups,
-    epsilon=1e-05,
-    weight=None,
-    bias=None,
-    data_format='NCHW',
-    name=None,
-):
+    x: Tensor,
+    num_groups: int,
+    epsilon: float = 1e-05,
+    weight: Tensor | None = None,
+    bias: Tensor | None = None,
+    data_format: DataLayout1D | DataLayout2D | DataLayout3D = 'NCHW',
+    name: str | None = None,
+) -> Tensor:
     """
     nn.GroupNorm is recommended.
     For more information, please refer to :ref:`api_paddle_nn_GroupNorm` .
@@ -663,7 +692,7 @@ def group_norm(
         bias(Tensor, optional): The bias Tensor of group_norm, with shape: attr:`[num_channels]`.
             Default: None.
         data_format(str, optional): Specify the input data format. Support "NCL", "NCHW", "NCDHW", "NLC", "NHWC" or "NDHWC". Default: "NCHW".
-        name(str, optional): Name for the GroupNorm, default is None. For more information, please refer to :ref:`api_guide_Name`..
+        name(str|None, optional): Name for the GroupNorm, default is None. For more information, please refer to :ref:`api_guide_Name`..
 
     Returns:
         Tensor, the output has the same shape with ``x``.
diff --git a/python/paddle/nn/functional/vision.py b/python/paddle/nn/functional/vision.py
@@ -18,7 +18,6 @@
 
 import paddle
 from paddle import _C_ops, _legacy_C_ops, in_dynamic_mode
-from paddle._typing import DataLayout2D, ShapeLike
 from paddle.base.framework import (
     in_dygraph_mode,
     in_dynamic_or_pir_mode,
@@ -32,6 +31,7 @@
 
 if TYPE_CHECKING:
     from paddle import Tensor
+    from paddle._typing import DataLayout2D, ShapeLike
 
 __all__ = []
 
diff --git a/python/paddle/nn/initializer/__init__.py b/python/paddle/nn/initializer/__init__.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# TODO: define the initializers to create a Parameter in neural network
+
 from ...base.initializer import set_global_initializer
 from .assign import (
     Assign,
diff --git a/python/paddle/nn/initializer/initializer.py b/python/paddle/nn/initializer/initializer.py
@@ -21,8 +21,6 @@
 import numpy as np
 from typing_extensions import TypeAlias
 
-import paddle
-
 from ...base.framework import (
     EagerParamBase,
     default_main_program,
@@ -31,6 +29,8 @@
 from .lazy_init import lazy_init_helper
 
 if TYPE_CHECKING:
+    import paddle
+
     _NonLinearity: TypeAlias = Literal[  # noqa: PYI047
         "sigmoid",
         "linear",
diff --git a/python/paddle/nn/layer/activation.py b/python/paddle/nn/layer/activation.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# TODO: define activation functions of neural network
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING
diff --git a/python/paddle/nn/layer/container.py b/python/paddle/nn/layer/container.py
@@ -21,12 +21,13 @@
 
 from typing_extensions import Self
 
-from paddle import Tensor
-
 from ...base.dygraph.base import param_guard
 from ...base.framework import Parameter
 from .layers import Layer
 
+if typing.TYPE_CHECKING:
+    from paddle import Tensor
+
 __all__ = []
 
 
diff --git a/python/paddle/nn/layer/distance.py b/python/paddle/nn/layer/distance.py
@@ -14,11 +14,14 @@
 
 from __future__ import annotations
 
-import paddle
+from typing import TYPE_CHECKING
 
 from .. import functional as F
 from .layers import Layer
 
+if TYPE_CHECKING:
+    import paddle
+
 __all__ = []
 
 
diff --git a/python/paddle/nn/layer/layers.py b/python/paddle/nn/layer/layers.py
@@ -27,7 +27,6 @@
 
 import paddle
 from paddle import Tensor, nn, profiler
-from paddle._typing import DTypeLike, PlaceLike, ShapeLike
 from paddle.autograd.backward_utils import ValueSet
 from paddle.base import core, framework, unique_name
 from paddle.base.core import VarDesc
@@ -56,7 +55,7 @@
 from paddle.utils import deprecated
 
 if TYPE_CHECKING:
-    from paddle._typing import ParamAttrLike
+    from paddle._typing import DTypeLike, ParamAttrLike, PlaceLike, ShapeLike
     from paddle.nn.initializer import Initializer
 
 
diff --git a/python/paddle/nn/layer/norm.py b/python/paddle/nn/layer/norm.py
diff --git a/python/paddle/nn/layer/pooling.py b/python/paddle/nn/layer/pooling.py
diff --git a/python/paddle/nn/layer/rnn.py b/python/paddle/nn/layer/rnn.py
diff --git a/python/paddle/nn/layer/transformer.py b/python/paddle/nn/layer/transformer.py
diff --git a/python/paddle/nn/layer/vision.py b/python/paddle/nn/layer/vision.py
