from __future__ import annotations

from typing import Any, Optional

OP_NAME_SCOPE: str
CLIP_OP_NAME_SCOPE: str
STEP_COUNTER: str
OP_ROLE_VAR_ATTR_NAME: Any
RPC_OP_ROLE_ATTR_NAME: Any
RPC_OP_ROLE_ATTR_VALUE: Any
LR_SCHED_OP_ROLE_ATTR_VALUE: Any
OPT_OP_ROLE_ATTR_VALUE: Any
op_role_attr_name: Any
SPARSE_OP_TYPE_DICT: Any
SPARSE_GRAD_OP_TYPE_DICT: Any
DEVICE_LIST: Any
COMMUNICATE_OPS_TYPE: Any
DEFAULT_DEVICE: str

def delete_optimizer_pass(program: Any, config: Any): ...
def distributed_ops_pass(program: Any, config: Any, use_ps_gpu: bool = ...): ...
def append_send_ops_pass(program: Any, config: Any): ...
def init_from_server_pass(program: Any, config: Any): ...
def fake_init_ops_pass(program: Any, config: Any): ...
def ps_gpu_pass(program: Any): ...
def delete_extra_optimizes_pass(program: Any, config: Any): ...
def find_heter_ops(program: Any, default_device: str = ...): ...
def create_heter_program(
    program: Any,
    config: Any,
    heter_program: Any,
    program_block_ops_list: Any,
    heter_ops: Any,
    block_var_detail: Any,
    current_device: Any,
    stage_id: Any,
) -> None: ...
def check_heter_compile_time_strategy(program: Any, config: Any, send_grad_var_list: Any) -> None: ...
def create_trainer_program(
    program: Any, origin_program: Any, config: Any, program_block_ops_list: Any, block_var_detail: Any
) -> None: ...
def insert_communicate_op(
    orign_program: Any,
    config: Any,
    heter_block: Any,
    stage_id: Any,
    first_op_index: Any,
    block_var_detail: Any,
    device: Any,
    is_forward: bool = ...,
): ...
def create_backward_block(program: Any, origin_program: Any, config: Any, bp_ops_list: Any, block_var_detail: Any): ...
def replace_ops_by_communicate_op(
    program: Any, config: Any, heter_block_index: Any, ops_list: Any, block_var_detail: Any
): ...
def remove_trainer_send_op(program: Any, config: Any, heter_block_index: Any, block_var_detail: Any) -> None: ...
def add_heter_send_op(program: Any, heter_program: Any, block: Any, block_var_detail: Any): ...
def find_send_op(program: Any): ...
def get_communicate_var_info(program: Any, block_index: Any, entrance_var_list: Any, type: 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_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 find_need_var_from_previous_block(
    need_add_vars: Any, block_var_detail: Any, current_index: Any, heter_ops: Any
): ...
def delete_block_useless_exit(program: Any, program_block_ops_list: Any, block_var_detail: Any): ...
def check_op_device(block: Any, device: Any) -> None: ...
def screen_persistables(program: Any, var_list: Any): ...
def insert_reshape_op(
    program: Any, block: Any, index: Any, var_name: Any, new_var_name: Any, new_var_shape: Optional[Any] = ...
) -> None: ...
def insert_send_concat_op(
    program: Any, block: Any, index: Any, var_name_list: Any, new_var_name: Any, new_var_shape: Any
) -> None: ...
def insert_recv_slice_op(
    program: Any,
    block: Any,
    index: Any,
    var_name: Any,
    var_shape: Any,
    dtype: Any,
    type: Any,
    new_var_name_list: Any,
    new_var_shape_list: Any,
) -> None: ...
def add_heter_trainer_useful_vars(
    config: Any, program: Any, heter_program: Any, heter_block: Any, static_var: Any
) -> None: ...
def delete_trainer_useless_var(config: Any, program: Any, static_var: Any): ...
def block_append_op(program: Any, origin_program: Any, block: Any, op: Any): ...
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 find_ops_list_input_output(program: Any, ops_list: Any): ...
def find_op_input_output(program: Any, block: Any, op: Any): ...
def get_vars_name_in_block(block: Any): ...
def is_same_op(op1: Any, op2: Any): ...
def delete_same_ops(block: Any, ops: Any) -> None: ...
