from __future__ import annotations

from typing import Any, Optional

def is_valid_list_index(list: Any, index: Any): ...
def is_dim_shard(mapping: Any): ...
def is_dim_replicate(mapping: Any): ...
def compute_compatible_dim_mapping(dim_mappings: Any): ...
def compute_compatible_dims_mapping(dims_mapping_list: Any): ...
def compute_compatible_process_mesh(process_mesh_list: Any): ...
def compute_compatible_and_update_dim_mapping(dims_mapping_list: Any, index_list: Any): ...
def append_distributed_attr_suffix(name: Any): ...
def remove_distributed_attr_suffix(name: Any): ...
def check_distributed_attr_for_program(program: Any, dist_context: Optional[Any] = ...): ...
def print_program_with_dist_attr(program: Any, dist_context: Optional[Any] = ...) -> None: ...
def make_data_unshard(dist_main_prog: Any, dist_startup_prog: Any, dist_context: Optional[Any] = ...) -> None: ...
def save_distributed_checkpoint(
    program: Any,
    checkpoint_path: Any,
    dist_attr_path: Any,
    addition_info: Optional[Any] = ...,
    is_integrated: bool = ...,
    dist_context: Optional[Any] = ...,
) -> None: ...
def load_distributed_checkpoint(checkpoint_path: Any, dist_attr_path: Any): ...
def load_checkpoint_into_program(
    checkpoint_path: Any, dist_attr_path: Any, program: Any, dist_context: Optional[Any] = ...
): ...
def load_parameter_into_program(param_dict: Any, program: Any) -> None: ...
def get_dist_attr(program: Any, dist_context: Optional[Any] = ...): ...
def merge_and_slice_parameter(dist_param_dict: Any, pre_dist_attr: Any, cur_dist_attr: Any): ...
def set_grad_var_shape(program: Any, dist_context: Any) -> None: ...

OP_ROLE_KEY: Any

def is_forward_op(op: Any): ...
def is_backward_op(op: Any): ...
def is_loss_op(op: Any): ...
def get_loss_op(block: Any): ...
def set_var_dist_attr(dist_context: Any, var: Any, dims_mapping: Any, process_mesh: Any, **kwargs: Any): ...
def naive_set_dist_op_attr_for_program_by_mesh_and_mapping(
    new_op: Any, process_mesh: Any, ref_mapping: Any, ctx: Any
) -> None: ...
def update_op_dims_mapping_by_default_dist_impl(dist_op: Any): ...
def update_op_dims_mapping_by_elementwise_like_dist_impl(dist_op: Any): ...
def get_all_distributed_main_program(serial_program_info: Any, dist_context: Any, parallelizer: Any): ...

class SerialProgramInfo:
    def __init__(
        self, train_program: Any, satrtup_program: Any, loss: Any, optimizer: Any, cluster: Optional[Any] = ...
    ) -> None: ...
    @property
    def train_program(self): ...
    @property
    def startup_program(self): ...
    @property
    def loss(self): ...
    @property
    def optimizer(self): ...
    @property
    def cluster(self): ...

def get_standalone_cost_data(distributed_programs: Any): ...
def set_dist_op_desc_original_id(dist_op_desc: Any, op_desc: Any, dist_context: Any) -> None: ...
def to_list(value: Any): ...
