[Typing] 增加类型 ParamAttrLike 并修改相应标注 (#65547)

megemini · web-flow · commit 418b6c7b82ed · 2024-06-28T20:26:51.000+08:00
diff --git a/python/paddle/_typing/__init__.py b/python/paddle/_typing/__init__.py
@@ -24,6 +24,7 @@
     NestedStructure as NestedStructure,
     Numberic as Numberic,
     NumbericSequence as NumbericSequence,
+    ParamAttrLike as ParamAttrLike,
     TensorIndex as TensorIndex,
     TensorLike as TensorLike,
     TensorOrTensors as TensorOrTensors,
diff --git a/python/paddle/_typing/basic.py b/python/paddle/_typing/basic.py
@@ -31,7 +31,9 @@
 from .backport import EllipsisType
 
 if TYPE_CHECKING:
-    from paddle import Tensor
+    from paddle import ParamAttr, Tensor
+    from paddle.nn.initializer import Initializer
+    from paddle.regularizer import WeightDecayRegularizer
 
 
 Numberic: TypeAlias = Union[int, float, complex, np.number, "Tensor"]
@@ -57,3 +59,7 @@
 NumbericSequence = Sequence[Numberic]
 NestedNumbericSequence: TypeAlias = NestedSequence[Numberic]
 TensorOrTensors: TypeAlias = Union["Tensor", Sequence["Tensor"]]
+
+ParamAttrLike: TypeAlias = Union[
+    "ParamAttr", "Initializer", "WeightDecayRegularizer", str, bool
+]
diff --git a/python/paddle/nn/layer/activation.py b/python/paddle/nn/layer/activation.py
@@ -25,8 +25,7 @@
 
 if TYPE_CHECKING:
     from paddle import Tensor
-    from paddle._typing import DataLayoutND
-    from paddle.base import ParamAttr
+    from paddle._typing import DataLayoutND, ParamAttrLike
 __all__ = []
 
 
@@ -489,7 +488,7 @@ def __init__(
         self,
         num_parameters: int = 1,
         init: float = 0.25,
-        weight_attr: ParamAttr | None = None,
+        weight_attr: ParamAttrLike | None = None,
         data_format: DataLayoutND = "NCHW",
         name: str | None = None,
     ) -> None:
diff --git a/python/paddle/nn/layer/common.py b/python/paddle/nn/layer/common.py
@@ -23,13 +23,14 @@
 from .layers import Layer
 
 if TYPE_CHECKING:
-    from paddle import ParamAttr, Tensor
+    from paddle import Tensor
     from paddle._typing import (
         DataLayout1D,
         DataLayout1DVariant,
         DataLayout2D,
         DataLayout3D,
         IntSequence,
+        ParamAttrLike,
         ShapeLike,
         Size2,
         Size4,
@@ -197,8 +198,8 @@ def __init__(
         self,
         in_features: int,
         out_features: int,
-        weight_attr: ParamAttr | None = None,
-        bias_attr: ParamAttr | bool | None = None,
+        weight_attr: ParamAttrLike | None = None,
+        bias_attr: ParamAttrLike | None = None,
         name: str | None = None,
     ) -> None:
         super().__init__()
@@ -744,8 +745,8 @@ def __init__(
         in1_features: int,
         in2_features: int,
         out_features: int,
-        weight_attr: ParamAttr | None = None,
-        bias_attr: ParamAttr | bool | None = None,
+        weight_attr: ParamAttrLike | None = None,
+        bias_attr: ParamAttrLike | None = None,
         name: str | None = None,
     ) -> None:
         super().__init__()
@@ -1765,7 +1766,7 @@ def __init__(
         embedding_dim: int,
         padding_idx: float | None = None,
         sparse: bool = False,
-        weight_attr: ParamAttr | None = None,
+        weight_attr: ParamAttrLike | None = None,
         name: str | None = None,
     ) -> None:
         super().__init__()
diff --git a/python/paddle/nn/layer/layers.py b/python/paddle/nn/layer/layers.py
@@ -56,6 +56,7 @@
 from paddle.utils import deprecated
 
 if TYPE_CHECKING:
+    from paddle._typing import ParamAttrLike
     from paddle.nn.initializer import Initializer
 
 
@@ -755,7 +756,7 @@ def register_forward_pre_hook(
     def create_parameter(
         self,
         shape: ShapeLike,
-        attr: ParamAttr | None = None,
+        attr: ParamAttrLike | None = None,
         dtype: DTypeLike | None = None,
         is_bias: bool = False,
         default_initializer: Initializer | None = None,
diff --git a/python/paddle/nn/layer/loss.py b/python/paddle/nn/layer/loss.py
@@ -26,7 +26,7 @@
 
 if TYPE_CHECKING:
     from paddle import Tensor
-    from paddle.base.param_attr import ParamAttr
+    from paddle._typing import ParamAttrLike
 
     from ..functional.loss import _ReduceMode
 
@@ -535,8 +535,8 @@ def __init__(
         self,
         feature_size: int,
         num_classes: int,
-        weight_attr: ParamAttr | None = None,
-        bias_attr: ParamAttr | bool | None = None,
+        weight_attr: ParamAttrLike | None = None,
+        bias_attr: ParamAttrLike | None = None,
         is_custom: bool = False,
         is_sparse: bool = False,
         name: str | None = None,
@@ -2491,8 +2491,8 @@ class AdaptiveLogSoftmaxWithLoss(Layer):
     in_features: int
     n_classes: int
     cutoffs: Sequence[int]
-    weight_attr: ParamAttr | None
-    bias_attr: ParamAttr | bool | None
+    weight_attr: ParamAttrLike | None
+    bias_attr: ParamAttrLike | None
     div_value: float
     head_bias: Tensor | None
     head_weight: Tensor
@@ -2504,8 +2504,8 @@ def __init__(
         in_features: int,
         n_classes: int,
         cutoffs: Sequence[int],
-        weight_attr: ParamAttr | None = None,
-        bias_attr: ParamAttr | bool | None = None,
+        weight_attr: ParamAttrLike | None = None,
+        bias_attr: ParamAttrLike | None = None,
         div_value: float = 4.0,
         head_bias: bool = False,
         name: str | None = None,
diff --git a/python/paddle/nn/layer/rnn.py b/python/paddle/nn/layer/rnn.py
@@ -51,7 +51,8 @@
 if TYPE_CHECKING:
     from typing import Literal
 
-    from paddle import ParamAttr, Tensor
+    from paddle import Tensor
+    from paddle._typing import ParamAttrLike
 
     _DirectionType = Literal["forward", "bidirect", "bidirectional"]
     _RNNType = Literal["LSTM", "GRU", "RNN_RELU", "RNN_TANH"]
@@ -808,10 +809,10 @@ def __init__(
         input_size: int,
         hidden_size: int,
         activation: _ActivationType | str = "tanh",
-        weight_ih_attr: ParamAttr | None = None,
-        weight_hh_attr: ParamAttr | None = None,
-        bias_ih_attr: ParamAttr | None = None,
-        bias_hh_attr: ParamAttr | None = None,
+        weight_ih_attr: ParamAttrLike | None = None,
+        weight_hh_attr: ParamAttrLike | None = None,
+        bias_ih_attr: ParamAttrLike | None = None,
+        bias_hh_attr: ParamAttrLike | None = None,
         name: str | None = None,
     ) -> None:
         super().__init__()
@@ -1012,10 +1013,10 @@ def __init__(
         self,
         input_size: int,
         hidden_size: int,
-        weight_ih_attr: ParamAttr | None = None,
-        weight_hh_attr: ParamAttr | None = None,
-        bias_ih_attr: ParamAttr | None = None,
-        bias_hh_attr: ParamAttr | None = None,
+        weight_ih_attr: ParamAttrLike | None = None,
+        weight_hh_attr: ParamAttrLike | None = None,
+        bias_ih_attr: ParamAttrLike | None = None,
+        bias_hh_attr: ParamAttrLike | None = None,
         proj_size: int = 0,
         name: str | None = None,
     ) -> None:
@@ -1222,10 +1223,10 @@ def __init__(
         self,
         input_size: int,
         hidden_size: int,
-        weight_ih_attr: ParamAttr | None = None,
-        weight_hh_attr: ParamAttr | None = None,
-        bias_ih_attr: ParamAttr | None = None,
-        bias_hh_attr: ParamAttr | None = None,
+        weight_ih_attr: ParamAttrLike | None = None,
+        weight_hh_attr: ParamAttrLike | None = None,
+        bias_ih_attr: ParamAttrLike | None = None,
+        bias_hh_attr: ParamAttrLike | None = None,
         name: str | None = None,
     ) -> None:
         super().__init__()
@@ -1525,10 +1526,10 @@ def __init__(
         direction: _DirectionType | str = "forward",
         time_major: bool = False,
         dropout: float = 0.0,
-        weight_ih_attr: ParamAttr | None = None,
-        weight_hh_attr: ParamAttr | None = None,
-        bias_ih_attr: ParamAttr | None = None,
-        bias_hh_attr: ParamAttr | None = None,
+        weight_ih_attr: ParamAttrLike | None = None,
+        weight_hh_attr: ParamAttrLike | None = None,
+        bias_ih_attr: ParamAttrLike | None = None,
+        bias_hh_attr: ParamAttrLike | None = None,
         proj_size: int = 0,
     ) -> None:
         super().__init__()
@@ -1939,10 +1940,10 @@ def __init__(
         time_major: bool = False,
         dropout: float = 0.0,
         activation: _ActivationType | str = "tanh",
-        weight_ih_attr: ParamAttr | None = None,
-        weight_hh_attr: ParamAttr | None = None,
-        bias_ih_attr: ParamAttr | None = None,
-        bias_hh_attr: ParamAttr | None = None,
+        weight_ih_attr: ParamAttrLike | None = None,
+        weight_hh_attr: ParamAttrLike | None = None,
+        bias_ih_attr: ParamAttrLike | None = None,
+        bias_hh_attr: ParamAttrLike | None = None,
         name: str | None = None,
     ) -> None:
         if activation == "tanh":
@@ -2082,10 +2083,10 @@ def __init__(
         direction: _DirectionType | str = "forward",
         time_major: bool = False,
         dropout: float = 0.0,
-        weight_ih_attr: ParamAttr | None = None,
-        weight_hh_attr: ParamAttr | None = None,
-        bias_ih_attr: ParamAttr | None = None,
-        bias_hh_attr: ParamAttr | None = None,
+        weight_ih_attr: ParamAttrLike | None = None,
+        weight_hh_attr: ParamAttrLike | None = None,
+        bias_ih_attr: ParamAttrLike | None = None,
+        bias_hh_attr: ParamAttrLike | None = None,
         proj_size: int = 0,
         name: str | None = None,
     ) -> None:
@@ -2203,10 +2204,10 @@ def __init__(
         direction: _DirectionType | str = "forward",
         time_major: bool = False,
         dropout: float = 0.0,
-        weight_ih_attr: ParamAttr | None = None,
-        weight_hh_attr: ParamAttr | None = None,
-        bias_ih_attr: ParamAttr | None = None,
-        bias_hh_attr: ParamAttr | None = None,
+        weight_ih_attr: ParamAttrLike | None = None,
+        weight_hh_attr: ParamAttrLike | None = None,
+        bias_ih_attr: ParamAttrLike | None = None,
+        bias_hh_attr: ParamAttrLike | None = None,
         name: str | None = None,
     ) -> None:
         super().__init__(
diff --git a/python/paddle/nn/layer/transformer.py b/python/paddle/nn/layer/transformer.py
@@ -34,6 +34,7 @@
 
 if TYPE_CHECKING:
     from paddle import Tensor
+    from paddle._typing import ParamAttrLike
 
 __all__ = []
 
@@ -47,7 +48,7 @@ def _convert_param_attr_to_list(
 
 @overload
 def _convert_param_attr_to_list(
-    param_attr: Sequence[ParamAttr] | ParamAttr | Sequence[bool] | bool | None,
+    param_attr: Sequence[ParamAttrLike] | ParamAttrLike | None,
     n: int,
 ) -> list[ParamAttr]:
     ...
@@ -192,8 +193,8 @@ def __init__(
         kdim: int | None = None,
         vdim: int | None = None,
         need_weights: bool = False,
-        weight_attr: ParamAttr | None = None,
-        bias_attr: ParamAttr | bool | None = None,
+        weight_attr: ParamAttrLike | None = None,
+        bias_attr: ParamAttrLike | None = None,
     ) -> None:
         super().__init__()
 
@@ -648,8 +649,8 @@ def __init__(
         attn_dropout: float | None = None,
         act_dropout: float | None = None,
         normalize_before: bool = False,
-        weight_attr: Sequence[ParamAttr] | ParamAttr | None = None,
-        bias_attr: ParamAttr | Sequence[ParamAttr] | bool | None = None,
+        weight_attr: ParamAttrLike | Sequence[ParamAttrLike] | None = None,
+        bias_attr: ParamAttrLike | Sequence[ParamAttrLike] | None = None,
         layer_norm_eps: float = 1e-5,
     ) -> None:
         self._config = locals()
@@ -1023,8 +1024,8 @@ def __init__(
         attn_dropout: float | None = None,
         act_dropout: float | None = None,
         normalize_before: bool = False,
-        weight_attr: ParamAttr | Sequence[ParamAttr] | None = None,
-        bias_attr: ParamAttr | Sequence[ParamAttr] | bool | None = None,
+        weight_attr: ParamAttrLike | Sequence[ParamAttrLike] | None = None,
+        bias_attr: ParamAttrLike | Sequence[ParamAttrLike] | None = None,
         layer_norm_eps: float = 1e-5,
     ) -> None:
         self._config = locals()
@@ -1563,8 +1564,8 @@ def __init__(
         attn_dropout: float | None = None,
         act_dropout: float | None = None,
         normalize_before: bool = False,
-        weight_attr: ParamAttr | Sequence[ParamAttr] | None = None,
-        bias_attr: ParamAttr | Sequence[ParamAttr] | bool | None = None,
+        weight_attr: ParamAttrLike | Sequence[ParamAttrLike] | None = None,
+        bias_attr: ParamAttrLike | Sequence[ParamAttrLike] | None = None,
         custom_encoder: Layer | None = None,
         custom_decoder: Layer | None = None,
     ) -> None:
diff --git a/python/paddle/tensor/creation.py b/python/paddle/tensor/creation.py
@@ -16,7 +16,7 @@
 
 import math
 import re
-from typing import Any, Sequence, overload
+from typing import TYPE_CHECKING, Any, Sequence, overload
 
 import numpy as np
 import numpy.typing as npt
@@ -55,6 +55,9 @@
     in_pir_mode,
 )
 
+if TYPE_CHECKING:
+    from paddle._typing import ParamAttrLike
+
 __all__ = []
 
 
@@ -174,7 +177,7 @@ def create_parameter(
     shape: ShapeLike,
     dtype: DTypeLike,
     name: str | None = None,
-    attr: ParamAttr | None = None,
+    attr: ParamAttrLike | None = None,
     is_bias: bool = False,
     default_initializer: paddle.nn.initializer.Initializer | None = None,
 ) -> paddle.Tensor:
