from __future__ import annotations

from typing import Any, Optional

from paddle.fluid.dygraph import Layer as Layer
from paddle.fluid.dygraph import LayerList as LayerList

from ..base.private_helper_function import wait_server_ready as wait_server_ready
from .common import OP_ROLE_KEY as OP_ROLE_KEY
from .common import OP_ROLE_VAR_KEY as OP_ROLE_VAR_KEY
from .common import CollectiveHelper as CollectiveHelper
from .common import OpRole as OpRole
from .common import is_backward_op as is_backward_op
from .common import is_loss_grad_op as is_loss_grad_op
from .common import is_optimizer_op as is_optimizer_op
from .meta_optimizer_base import MetaOptimizerBase as MetaOptimizerBase

class RawProgramOptimizer(MetaOptimizerBase):
    inner_opt: Any = ...
    meta_optimizers_white_list: Any = ...
    meta_optimizers_black_list: Any = ...
    global_ring_id: int = ...
    def __init__(self, optimizer: Any) -> None: ...
    endpoints: Any = ...
    current_endpoint: Any = ...
    rank: Any = ...
    nranks: Any = ...
    startup_program: Any = ...
    main_program: Any = ...
    def minimize_impl(
        self,
        loss: Any,
        startup_program: Optional[Any] = ...,
        parameter_list: Optional[Any] = ...,
        no_grad_set: Optional[Any] = ...,
    ): ...
