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

python/paddle/incubate/optimizer/lbfgs.py

@@ -11,17 +11,37 @@
 # 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, 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 typing_extensions import NotRequired, TypedDict
+
+    from paddle import Tensor
+    from paddle.nn.clip import GradientClipBase
+    from paddle.optimizer.lr import LRScheduler
+    from paddle.regularizer import WeightDecayRegularizer
+
+    class _ParameterConfig(TypedDict):
+        params: Sequence[Tensor]
+        weight_decay: NotRequired[float | WeightDecayRegularizer | None]
+        learning_rate: NotRequired[float | Tensor | LRScheduler | None]
+
+    _T = TypeVar('_T', covariant=True)
+    _KT = TypeVar("_KT")
+    _VT = TypeVar("_VT")
+
 
 @deprecated(since="2.5.0", update_to="paddle.optimizer.LBFGS", level=1)
 class LBFGS(Optimizer):
@@ -116,20 +136,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: str | None
+    state: dict[_KT, _VT]
+
     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: str | 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 +194,7 @@ def __init__(
 
         self._numel_cache = None
 
-    def state_dict(self):
+    def state_dict(self) -> dict[str, Tensor]:
         r"""Returns the state of the optimizer as a :class:`dict`.
 
         Return:
@@ -226,7 +255,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]) -> _T:
         """
         Performs a single optimization step.
 

python/paddle/incubate/optimizer/lookahead.py

@@ -11,16 +11,24 @@
 # 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
 from paddle.base.dygraph import base as imperative_base
 from paddle.base.framework import Variable
-from paddle.base.layer_helper import LayerHelper
+from paddle.base.layer_helper import LayerHelper, LayerHelperBase
 from paddle.framework import in_pir_mode
 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, Program
+
+
 __all__ = []
 
 
@@ -112,9 +120,21 @@ class LookAhead(Optimizer):
 
     """
 
+    inner_optimizer: Optimizer
+    alpha: float
+    k: int
+    type: str
+    helper: LayerHelperBase | None
+
     _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 +172,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 +292,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``.
 

python/paddle/incubate/optimizer/modelaverage.py

@@ -11,12 +11,15 @@
 # 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
 from paddle.base import framework
 from paddle.base.dygraph import base as imperative_base
-from paddle.base.layer_helper import LayerHelper
+from paddle.base.layer_helper import LayerHelper, LayerHelperBase
 from paddle.base.wrapped_decorator import signature_safe_contextmanager
 from paddle.framework import (
     in_dynamic_mode,
@@ -25,6 +28,23 @@
 )
 from paddle.optimizer import Optimizer
 
+if TYPE_CHECKING:
+    from collections.abc import Generator, Sequence
+
+    from typing_extensions import NotRequired, TypedDict
+
+    from paddle import Tensor
+    from paddle.base.framework import Program
+    from paddle.optimizer.lr import LRScheduler
+    from paddle.regularizer import WeightDecayRegularizer
+    from paddle.static import Executor
+
+    class _ParameterConfig(TypedDict):
+        params: Sequence[Tensor]
+        weight_decay: NotRequired[float | WeightDecayRegularizer | None]
+        learning_rate: NotRequired[float | Tensor | LRScheduler | None]
+
+
 __all__ = []
 
 
@@ -169,14 +189,22 @@ class ModelAverage(Optimizer):
 
     """
 
+    helper: LayerHelperBase | None
+    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 +324,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 +382,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 +427,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 +508,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.