from __future__ import annotations

from typing import Any, Optional

from .common import OP_ROLE_KEY as OP_ROLE_KEY
from .common import CollectiveHelper as CollectiveHelper
from .common import OpRole as OpRole
from .common import is_update_op as is_update_op
from .meta_optimizer_base import MetaOptimizerBase as MetaOptimizerBase

class LocalSGDOptimizer(MetaOptimizerBase):
    inner_opt: Any = ...
    meta_optimizers_white_list: Any = ...
    meta_optimizers_black_list: Any = ...
    snapshot_key: str = ...
    def __init__(self, optimizer: Any) -> None: ...
    def snapshot_name(self, param_name: Any): ...
    def create_snapshot_vars(self, program: Any): ...
    def init_snapshot_vars(self, startup_program: Any, param2snapshot: Any) -> None: ...
    nrings: int = ...
    def minimize_impl(
        self,
        loss: Any,
        startup_program: Optional[Any] = ...,
        parameter_list: Optional[Any] = ...,
        no_grad_set: Optional[Any] = ...,
    ): ...

class AdaptiveLocalSGDOptimizer(MetaOptimizerBase):
    inner_opt: Any = ...
    meta_optimizers_white_list: Any = ...
    meta_optimizers_black_list: Any = ...
    snapshot_key: str = ...
    def __init__(self, optimizer: Any) -> None: ...
    def snapshot_name(self, param_name: Any): ...
    def create_snapshot_vars(self, program: Any): ...
    def init_snapshot_vars(self, startup_program: Any, param2snapshot: Any) -> None: ...
    nrings: int = ...
    def minimize_impl(
        self,
        loss: Any,
        startup_program: Optional[Any] = ...,
        parameter_list: Optional[Any] = ...,
        no_grad_set: Optional[Any] = ...,
    ): ...
