[Typing][C-81,C-102,C-103,C-105] Add type annotations for {fused_dropout_add,bfgs,lbfgs,math}.py in incubate APIs (#67233)

Lans1ot · web-flow · commit a1d73ae72d79 · 2024-08-13T20:50:38.000+08:00
* 添加类型信息

* 追加类型信息

* 按要求进行修改

* 修改参数取值范围

* 修改默认参数

* 修改TYPE_CHECKING语句块位置
diff --git a/python/paddle/incubate/nn/layer/fused_dropout_add.py b/python/paddle/incubate/nn/layer/fused_dropout_add.py
@@ -12,9 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal
+
 from paddle.incubate.nn import functional as F
 from paddle.nn import Layer
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+
 
 class FusedDropoutAdd(Layer):
     r"""
@@ -57,13 +64,20 @@ class FusedDropoutAdd(Layer):
             >>> out = m(x, y)
     """
 
-    def __init__(self, p=0.5, mode="upscale_in_train", name=None):
+    def __init__(
+        self,
+        p: float = 0.5,
+        mode: Literal[
+            'upscale_in_train', 'downscale_in_infer'
+        ] = "upscale_in_train",
+        name: str | None = None,
+    ) -> None:
         super().__init__()
         self.p = p
         self.mode = mode
         self.name = name
 
-    def forward(self, x, y):
+    def forward(self, x: Tensor, y: Tensor) -> Tensor:
         out = F.fused_dropout_add(
             x,
             y,
@@ -74,6 +88,6 @@ def forward(self, x, y):
         )
         return out
 
-    def extra_repr(self):
+    def extra_repr(self) -> str:
         name_str = f', name={self.name}' if self.name else ''
         return f'p={self.p}, mode={self.mode}{name_str}'
diff --git a/python/paddle/incubate/optimizer/functional/bfgs.py b/python/paddle/incubate/optimizer/functional/bfgs.py
@@ -12,6 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal
+
 import numpy as np
 
 import paddle
@@ -23,20 +27,25 @@
     check_input_type,
 )
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from paddle import Tensor
+
 
 def minimize_bfgs(
-    objective_func,
-    initial_position,
-    max_iters=50,
-    tolerance_grad=1e-7,
-    tolerance_change=1e-9,
-    initial_inverse_hessian_estimate=None,
-    line_search_fn='strong_wolfe',
-    max_line_search_iters=50,
-    initial_step_length=1.0,
-    dtype='float32',
-    name=None,
-):
+    objective_func: Callable[[Tensor], Tensor],
+    initial_position: Tensor,
+    max_iters: int = 50,
+    tolerance_grad: float = 1e-7,
+    tolerance_change: float = 1e-9,
+    initial_inverse_hessian_estimate: Tensor | None = None,
+    line_search_fn: Literal['strong_wolfe'] = 'strong_wolfe',
+    max_line_search_iters: int = 50,
+    initial_step_length: float = 1.0,
+    dtype: Literal['float32', 'float64'] = 'float32',
+    name: str | None = None,
+) -> tuple[bool, int, Tensor, Tensor, Tensor, Tensor]:
     r"""
     Minimizes a differentiable function `func` using the BFGS method.
     The BFGS is a quasi-Newton method for solving an unconstrained optimization problem over a differentiable function.
diff --git a/python/paddle/incubate/optimizer/functional/lbfgs.py b/python/paddle/incubate/optimizer/functional/lbfgs.py
@@ -12,6 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal
+
 import numpy as np
 
 import paddle
@@ -23,21 +27,26 @@
     check_input_type,
 )
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from paddle import Tensor
+
 
 def minimize_lbfgs(
-    objective_func,
-    initial_position,
-    history_size=100,
-    max_iters=50,
-    tolerance_grad=1e-8,
-    tolerance_change=1e-8,
-    initial_inverse_hessian_estimate=None,
-    line_search_fn='strong_wolfe',
-    max_line_search_iters=50,
-    initial_step_length=1.0,
-    dtype='float32',
-    name=None,
-):
+    objective_func: Callable[[Tensor], Tensor],
+    initial_position: Tensor,
+    history_size: int = 100,
+    max_iters: int = 50,
+    tolerance_grad: float = 1e-8,
+    tolerance_change: float = 1e-8,
+    initial_inverse_hessian_estimate: Tensor | None = None,
+    line_search_fn: Literal['strong_wolfe'] = 'strong_wolfe',
+    max_line_search_iters: int = 50,
+    initial_step_length: int = 1.0,
+    dtype: Literal['float32', 'float64'] = 'float32',
+    name: str | None = None,
+) -> tuple[bool, int, Tensor, Tensor, Tensor]:
     r"""
     Minimizes a differentiable function `func` using the L-BFGS method.
     The L-BFGS is a quasi-Newton method for solving an unconstrained optimization problem over a differentiable function.
diff --git a/python/paddle/incubate/tensor/math.py b/python/paddle/incubate/tensor/math.py
@@ -12,12 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 from paddle import _C_ops
 from paddle.base.data_feeder import check_variable_and_dtype
 from paddle.base.layer_helper import LayerHelper
 from paddle.framework import in_dynamic_or_pir_mode
 from paddle.utils import deprecated
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+
 __all__ = []
 
 
@@ -27,7 +34,9 @@
     level=1,
     reason="paddle.incubate.segment_sum will be removed in future",
 )
-def segment_sum(data, segment_ids, name=None):
+def segment_sum(
+    data: Tensor, segment_ids: Tensor, name: str | None = None
+) -> Tensor:
     r"""
     Segment Sum Operator.
 
@@ -94,7 +103,9 @@ def segment_sum(data, segment_ids, name=None):
     level=1,
     reason="paddle.incubate.segment_mean will be removed in future",
 )
-def segment_mean(data, segment_ids, name=None):
+def segment_mean(
+    data: Tensor, segment_ids: Tensor, name: str | None = None
+) -> Tensor:
     r"""
     Segment Mean Operator.
 
@@ -163,7 +174,9 @@ def segment_mean(data, segment_ids, name=None):
     level=1,
     reason="paddle.incubate.segment_min will be removed in future",
 )
-def segment_min(data, segment_ids, name=None):
+def segment_min(
+    data: Tensor, segment_ids: Tensor, name: str | None = None
+) -> Tensor:
     r"""
     Segment min operator.
 
@@ -231,7 +244,9 @@ def segment_min(data, segment_ids, name=None):
     level=1,
     reason="paddle.incubate.segment_max will be removed in future",
 )
-def segment_max(data, segment_ids, name=None):
+def segment_max(
+    data: Tensor, segment_ids: Tensor, name: str | None = None
+) -> Tensor:
     r"""
     Segment max operator.
 
