from __future__ import annotations

from typing import Any, Optional

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
backward: Any
DEVICE_LIST: Any
COMMUNICATE_OPS_TYPE: Any
SPARSE_OP_LIST: Any
SPARSE_OP_TYPE_DICT: Any
SPARSE_GRAD_OP_TYPE_DICT: Any
DEFAULT_DEVICE: str
DATA_NORM_NAME: Any
DATA_NORM_GRAD_NAME: Any

def logger_config(log_path: Any, logging_name: Any): ...

ps_log_root_dir: str
logger: Any

class DistributedMode:
    SYNC: int = ...
    ASYNC: int = ...
    HALF_ASYNC: int = ...
    GEO: int = ...
    FL: int = ...

class TrainerRuntimeConfig:
    mode: Any = ...
    runtime_configs: Any = ...
    def __init__(self, valid_strategy: Any) -> None: ...
    def get_communicator_flags(self): ...

def get_lr_ops(program: Any): ...
def get_optimize_ops(_program: Any): ...
def get_dist_env(): ...
def get_role_id(role_maker: Any): ...
def get_ps_endpoint(role_maker: Any): ...
def get_ps_endpoints(role_maker: Any): ...
def get_heter_worker_endpoint(role_maker: Any): ...
def get_trainer_endpoint(role_maker: Any): ...
def get_previous_stage_trainers(role_maker: Any): ...
def is_distributed_sparse_op(op: Any): ...
def get_sparse_tablename(op: Any): ...
def is_sparse_op(op: Any): ...
def get_sparse_tablenames(programs: Any, is_distributed: Any): ...
def get_trainers(role_maker: Any): ...
def get_dense_send_context(
    program: Any, send_ctx: Any, idx: Any, merged_dense_pairs: Any, trainer_id: Any, split_dense_table: bool = ...
): ...
def get_geo_trainer_send_context(context: Any): ...
def get_the_one_send_context(
    context: Any, split_dense_table: bool = ..., use_origin_program: bool = ..., ep_list: Optional[Any] = ...
): ...
def find_heter_ops(program: Any, default_device: str = ...): ...
def union_forward_gradient_op(program_block_ops_list: Any): ...
def find_block_joints(program: Any, program_block_ops_list: Any, heter_ops: Any): ...
def find_ops_list_input_output(program: Any, ops_list: Any): ...
def find_entrance_exit_private(program: Any, program_block_ops_list: Any): ...
def entrance_exit_check(program: Any, program_block_ops_list: Any, block_var_detail: Any, heter_ops: Any): ...
def delete_block_useless_exit(program: Any, program_block_ops_list: Any, block_var_detail: Any): ...
def get_communicate_var_info(program: Any, block_index: Any, entrance_var_list: Any, type: str = ...): ...
def add_vars_by_var_list(var_name_list: Any, origin_program: Any, program: Any, block: Any) -> None: ...
def get_varlist_from_op_map(var_map: Any): ...
def screen_persistables(program: Any, var_list: Any): ...
def block_append_op(program: Any, origin_program: Any, block: Any, op: Any): ...
def get_next_stage_trainers(role_maker: Any): ...
def insert_communicate_op(
    orign_program: Any,
    role_maker: Any,
    heter_block: Any,
    stage_id: Any,
    first_op_index: Any,
    block_var_detail: Any,
    device: Any,
    is_forward: bool = ...,
): ...
def get_the_one_recv_context(
    context: Any, is_dense: bool = ..., split_dense_table: bool = ..., use_origin_program: bool = ...
): ...

dtype_to_size: Any

def get_var_mem_size(var: Any): ...

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

def build_var_distributed(context: Any) -> None: ...
def get_param_grads(origin_program: Any): ...
def delete_ops(block: Any, ops: Any) -> None: ...
def find_send_op(program: Any): ...
def find_op_input_output(program: Any, block: Any, op: Any): ...
def add_heter_send_op(program: Any, heter_program: Any, block: Any, block_var_detail: Any): ...
def get_vars_name_in_block(block: Any): ...
def delete_trainer_useless_var(program: Any, static_var: Any): ...
def create_backward_block(program: Any, origin_program: Any, bp_ops_list: Any, block_var_detail: Any): ...
def debug_program(file: Any, program: Any) -> None: ...
