from __future__ import annotations

from typing import Any, Optional

FLEET_GLOBAL_DICT: Any

class DistributedOptimizerImplBase:
    def __init__(self, optimizer: Any) -> None: ...
    def minimize(
        self,
        losses: Any,
        startup_program: Optional[Any] = ...,
        parameter_list: Optional[Any] = ...,
        no_grad_set: Optional[Any] = ...,
    ) -> None: ...

class DistributedAdam(DistributedOptimizerImplBase):
    type: str = ...
    data_norm_name: Any = ...
    supported_embedding_types: Any = ...
    supported_embedding_grad_types: Any = ...
    op_role_key: Any = ...
    def __init__(self, optimizer: Any) -> None: ...
