[Typing][C-95, C-96, C-99]Add annotations for python/paddle/incubate/optimizer/{lookahead, modelaverage, lbfgs}.py (#67448)

inaomIIsfarell · web-flow · commit 99daec499312 · 2024-08-21T14:56:16.000+08:00
diff --git a/python/paddle/incubate/optimizer/lbfgs.py b/python/paddle/incubate/optimizer/lbfgs.py
@@ -11,17 +11,28 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+from __future__ import annotations
 
 from collections import defaultdict
 from functools import reduce
+from typing import TYPE_CHECKING, Any, Literal, TypeVar
 
 import paddle
 from paddle.optimizer import Optimizer
 from paddle.utils import deprecated
 
 from .line_search_dygraph import _strong_wolfe
 
+if TYPE_CHECKING:
+    from collections.abc import Callable, Sequence
+
+    from paddle import Tensor
+    from paddle.nn.clip import GradientClipBase
+    from paddle.optimizer.optimizer import _ParameterConfig
+    from paddle.regularizer import WeightDecayRegularizer
+
+    _T_co = TypeVar('_T_co', covariant=True)
+
 
 @deprecated(since="2.5.0", update_to="paddle.optimizer.LBFGS", level=1)
 class LBFGS(Optimizer):
@@ -116,20 +127,29 @@ class LBFGS(Optimizer):
 
     """
 
+    learning_rate: float
+    max_iter: int
+    max_eval: int
+    tolerance_grad: float
+    tolerance_change: float
+    history_size: int
+    line_search_fn: Literal['strong_wolfe'] | None
+    state: dict[str, dict[str, Any]]
+
     def __init__(
         self,
-        learning_rate=1.0,
-        max_iter=20,
-        max_eval=None,
-        tolerance_grad=1e-7,
-        tolerance_change=1e-9,
-        history_size=100,
-        line_search_fn=None,
-        parameters=None,
-        weight_decay=None,
-        grad_clip=None,
-        name=None,
-    ):
+        learning_rate: float = 1.0,
+        max_iter: int = 20,
+        max_eval: int | None = None,
+        tolerance_grad: float = 1e-7,
+        tolerance_change: float = 1e-9,
+        history_size: int = 100,
+        line_search_fn: Literal['strong_wolfe'] | None = None,
+        parameters: Sequence[Tensor] | Sequence[_ParameterConfig] | None = None,
+        weight_decay: float | WeightDecayRegularizer | None = None,
+        grad_clip: GradientClipBase | None = None,
+        name: str | None = None,
+    ) -> Tensor:
         if max_eval is None:
             max_eval = max_iter * 5 // 4
 
@@ -165,7 +185,7 @@ def __init__(
 
         self._numel_cache = None
 
-    def state_dict(self):
+    def state_dict(self) -> dict[str, dict[str, Any]]:
         r"""Returns the state of the optimizer as a :class:`dict`.
 
         Return:
@@ -226,7 +246,7 @@ def _directional_evaluate(self, closure, x, alpha, d):
         self._set_param(x)
         return loss, flat_grad
 
-    def step(self, closure):
+    def step(self, closure: Callable[[], _T_co]) -> _T_co:
         """
         Performs a single optimization step.
 
diff --git a/python/paddle/incubate/optimizer/lookahead.py b/python/paddle/incubate/optimizer/lookahead.py
@@ -11,6 +11,9 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
 
 import paddle
 from paddle.base import framework, unique_name
@@ -21,6 +24,12 @@
 from paddle.optimizer import Optimizer
 from paddle.pir.core import create_parameter
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle.base.framework import Operator
+    from paddle.static import Program
+
+
 __all__ = []
 
 
@@ -112,9 +121,21 @@ class LookAhead(Optimizer):
 
     """
 
+    inner_optimizer: Optimizer
+    alpha: float
+    k: int
+    type: str
+    helper: LayerHelper
+
     _slow_str = "slow"
 
-    def __init__(self, inner_optimizer, alpha=0.5, k=5, name=None):
+    def __init__(
+        self,
+        inner_optimizer: Optimizer,
+        alpha: float = 0.5,
+        k: int = 5,
+        name: str | None = None,
+    ) -> None:
         assert inner_optimizer is not None, "inner optimizer can not be None"
         assert (
             0.0 <= alpha <= 1.0
@@ -152,7 +173,7 @@ def _set_auxiliary_var(self, key, val):
 
     @framework.dygraph_only
     @imperative_base.no_grad
-    def step(self):
+    def step(self) -> None:
         """
         Execute the optimizer and update parameters once.
 
@@ -272,8 +293,12 @@ def _append_optimize_op(self, block, param_and_grad):
 
     @imperative_base.no_grad
     def minimize(
-        self, loss, startup_program=None, parameters=None, no_grad_set=None
-    ):
+        self,
+        loss: Tensor,
+        startup_program: Program | None = None,
+        parameters: list[Tensor] | list[str] | None = None,
+        no_grad_set: set[Tensor] | set[str] | None = None,
+    ) -> tuple[list[Operator], list[tuple[Tensor, Tensor]]]:
         """
         Add operations to minimize ``loss`` by updating ``parameters``.
 
diff --git a/python/paddle/incubate/optimizer/modelaverage.py b/python/paddle/incubate/optimizer/modelaverage.py
@@ -11,6 +11,9 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
 
 import paddle
 from paddle import _C_ops
@@ -25,6 +28,14 @@
 )
 from paddle.optimizer import Optimizer
 
+if TYPE_CHECKING:
+    from collections.abc import Generator, Sequence
+
+    from paddle import Tensor
+    from paddle.optimizer.optimizer import _ParameterConfig
+    from paddle.static import Executor, Program
+
+
 __all__ = []
 
 
@@ -169,14 +180,22 @@ class ModelAverage(Optimizer):
 
     """
 
+    helper: LayerHelper
+    average_window: float
+    min_average_window: int
+    max_average_window: int
+    type: str
+    apply_program: Program
+    restore_program: Program
+
     def __init__(
         self,
-        average_window_rate,
-        parameters=None,
-        min_average_window=10000,
-        max_average_window=10000,
-        name=None,
-    ):
+        average_window_rate: float,
+        parameters: Sequence[Tensor] | Sequence[_ParameterConfig] | None = None,
+        min_average_window: int = 10000,
+        max_average_window: int = 10000,
+        name: str | None = None,
+    ) -> None:
         super().__init__(
             learning_rate=0.0,
             parameters=parameters,
@@ -296,8 +315,12 @@ def _append_optimize_op(self, block, param_and_grad):
 
     @imperative_base.no_grad
     def minimize(
-        self, loss, startup_program=None, parameters=None, no_grad_set=None
-    ):
+        self,
+        loss: Tensor,
+        startup_program: Program | None = None,
+        parameters: list[Tensor] | None = None,
+        no_grad_set: set[Tensor] | set[str] | None = None,
+    ) -> None:
         """
         Add operations to minimize ``loss`` by updating ``parameters``.
 
@@ -350,7 +373,7 @@ def minimize(
 
     @framework.dygraph_only
     @imperative_base.no_grad
-    def step(self):
+    def step(self) -> None:
         """
         Execute the optimizer and update parameters once.
 
@@ -395,7 +418,9 @@ def step(self):
 
     @signature_safe_contextmanager
     @imperative_base.no_grad
-    def apply(self, executor=None, need_restore=True):
+    def apply(
+        self, executor: Executor | None = None, need_restore: bool = True
+    ) -> Generator[None, None, None]:
         """
         Apply the average of the cumulative ``Parameter`` to the parameters of the current model.
 
@@ -474,7 +499,7 @@ def apply(self, executor=None, need_restore=True):
                 self.restore(executor)
 
     @imperative_base.no_grad
-    def restore(self, executor=None):
+    def restore(self, executor: Executor | None = None) -> None:
         """
         Restore ``Parameter`` values of current model.
 
