from __future__ import annotations

from typing import Any, Optional

from paddle.fluid.incubate.fleet.parameter_server.ir.ps_dispatcher import (
    PSDispatcher as PSDispatcher,
)
from paddle.fluid.transpiler.details.program_utils import delete_ops as delete_ops

OP_NAME_SCOPE: str
CLIP_OP_NAME_SCOPE: str
STEP_COUNTER: str
LEARNING_RATE_DECAY_COUNTER: str
OP_ROLE_VAR_ATTR_NAME: Any
RPC_OP_ROLE_ATTR_NAME: Any
RPC_OP_ROLE_ATTR_VALUE: Any
op_role_attr_name: Any
LR_SCHED_OP_ROLE_ATTR_VALUE: Any
OPT_OP_ROLE_ATTR_VALUE: Any
SPARSE_OP_LIST: Any
SPARSE_OP_TYPE_DICT: Any

def is_sparse_op(op: Any): ...
def is_distributed_sparse_op(op: Any): ...
def get_sparse_tablename(op: Any): ...
def get_sparse_tablenames(program: Any, is_distributed: Any): ...

class MergedVariable:
    merged_var: Any = ...
    ordered_vars: Any = ...
    offsets: Any = ...
    def __init__(self, merged: Any, ordered: Any, offsets: Any) -> None: ...

def Singleton(cls): ...

class CompileTimeStrategy:
    min_block_size: int = ...
    origin_main_program: Any = ...
    origin_startup_program: Any = ...
    origin_ps_main_program: Any = ...
    origin_ps_startup_program: Any = ...
    strategy: Any = ...
    role_maker: Any = ...
    use_ps_gpu: bool = ...
    is_heter_ps_mode: Any = ...
    origin_sparse_pairs: Any = ...
    origin_dense_pairs: Any = ...
    merged_variables_pairs: Any = ...
    merged_dense_pairs: Any = ...
    merged_sparse_pairs: Any = ...
    merged_variable_map: Any = ...
    param_name_to_grad_name: Any = ...
    grad_name_to_param_name: Any = ...
    param_grad_ep_mapping: Any = ...
    grad_param_mapping: Any = ...
    tensor_table_dict: Any = ...
    origin_merged_variables_pairs: Any = ...
    origin_merged_dense_pairs: Any = ...
    origin_merged_sparse_pairs: Any = ...
    def __init__(self, main_program: Any, startup_program: Any, strategy: Any, role_maker: Any) -> None: ...
    def get_distributed_mode(self): ...
    def is_sync_mode(self): ...
    def is_geo_mode(self): ...
    def is_async_mode(self): ...
    def get_role_id(self): ...
    def get_trainers(self): ...
    def get_ps_endpoint(self): ...
    def get_ps_endpoints(self): ...
    def get_heter_worker_endpoints(self): ...
    def get_next_stage_trainers(self): ...
    def get_heter_worker_endpoint(self): ...
    def get_trainer_endpoints(self): ...
    def get_trainer_endpoint(self): ...
    def get_previous_stage_trainers(self): ...
    def get_origin_programs(self): ...
    def get_origin_main_program(self): ...
    def get_origin_startup_program(self): ...
    def set_origin_ps_main_program(self, program: Any) -> None: ...
    def set_origin_ps_startup_program(self, program: Any) -> None: ...
    def get_origin_ps_main_program(self): ...
    def get_origin_ps_startup_program(self): ...
    def add_tensor_table(
        self,
        feed_var_name: Any,
        fetch_var_name: str = ...,
        startup_program: Optional[Any] = ...,
        main_program: Optional[Any] = ...,
        tensor_table_class: str = ...,
    ) -> None: ...
    def get_tensor_table_dict(self): ...
    def get_sparse_varname_on_ps(self, is_distributed: Any, endpoint: Optional[Any] = ...): ...
    def get_optimize_varname_on_ps(self, param_name: Any): ...
    def build_ctx(
        self, vars: Any, mapping: Any, is_grad: Any, is_sparse: Any, is_send: Any, is_distributed: bool = ...
    ): ...
    def get_trainer_send_context(self): ...
    def get_communicator_send_context(self): ...
    def get_communicator_recv_context(self, recv_type: int = ..., use_origin_program: bool = ...): ...
    def get_the_one_trainer_send_context(self, split_dense_table: Any): ...
    def get_dense_send_context(
        self, send_ctx: Any, idx: Any, merged_dense_pairs: Any, trainer_id: Any, split_dense_table: bool = ...
    ): ...
    def get_the_one_send_context(
        self, split_dense_table: bool = ..., use_origin_program: bool = ..., ep_list: Optional[Any] = ...
    ): ...
    def get_the_one_recv_context(
        self, is_dense: bool = ..., split_dense_table: bool = ..., use_origin_program: bool = ...
    ): ...
    def get_server_runtime_config(self): ...
    def get_var_distributed(self, varname: Any, is_param: Any): ...
    def get_param_grads(self): ...
    def remove_var_pair_by_grad(self, var_name: Any) -> None: ...
