from __future__ import annotations

from typing import Any, Optional

from .dist_attribute import OperatorDistributedAttribute as OperatorDistributedAttribute
from .dist_attribute import TensorDistributedAttribute as TensorDistributedAttribute
from .dist_context import DistributedContext as DistributedContext
from .process_group import ProcessGroup as ProcessGroup
from .process_group import new_process_group as new_process_group

def get_var_with_recursion(var_name: Any, block: Any, program: Any): ...

class AllGatherOpDesc:
    def __init__(self, group: Any) -> None: ...
    @property
    def group(self): ...
    @property
    def desc(self): ...

class SendOpDesc:
    def __init__(self, partition_index: Any, dst: Any) -> None: ...
    @property
    def partition_index(self): ...
    @property
    def dst(self): ...
    @property
    def desc(self): ...

class RecvOpDesc:
    def __init__(self, partition_index: Any, src: Any) -> None: ...
    @property
    def partition_index(self): ...
    @property
    def src(self): ...
    @property
    def desc(self): ...

class SliceOpDesc:
    def __init__(self, starts: Any, ends: Any, axes: Any) -> None: ...
    @property
    def starts(self): ...
    @property
    def ends(self): ...
    @property
    def axes(self): ...
    @property
    def desc(self): ...

class ConcatOpDesc:
    def __init__(self, partition_index_list: Any) -> None: ...
    @property
    def partition_index_list(self): ...
    @property
    def desc(self): ...

class Inserter:
    @staticmethod
    def insert_send_op(block: Any, idx: Any, tensor: Any, dst: Any, op_role: Any) -> None: ...
    @staticmethod
    def insert_recv_op(block: Any, idx: Any, tensor: Any, src: Any, op_role: Any) -> None: ...
    @staticmethod
    def insert_concat_op(block: Any, idx: Any, tensors: Any, axis: Any, op_role: Any): ...
    @staticmethod
    def insert_slice_op(
        block: Any, idx: Any, tensor: Any, starts: Any, ends: Any, axes: Any, new_var_name: Any, op_role: Any
    ): ...
    @staticmethod
    def insert_split_op(block: Any, idx: Any, tensor: Any, num_or_sections: Any, op_role: Any): ...
    @staticmethod
    def insert_fill_constant_op(block: Any, idx: Any, op_role: Any): ...
    @staticmethod
    def insert_allgather_op(block: Any, idx: Any, tensor: Any, ranks: Any, op_role: Any): ...
    @staticmethod
    def concat_partitions_with_op(
        partition_tensor_list: Any, tensor: Any, partition_index: Any, block: Any, idx: Any, op_role: Any
    ) -> None: ...

class Remover:
    @staticmethod
    def remove_no_need_ops(auto_parallel_main_prog: Any, dist_context: Any, rank_id: Any) -> None: ...
    @staticmethod
    def remove_no_need_vars(auto_parallel_main_prog: Any, dist_params_grads: Any) -> None: ...
    @staticmethod
    def remove_no_need_in_main(
        auto_parallel_main_prog: Any, dist_context: Any, rank_id: Any, dist_params_grads: Any
    ) -> None: ...
    @staticmethod
    def remove_no_need_in_startup(auto_parallel_main_prog: Any, auto_parallel_startup_prog: Any) -> None: ...

class Resharder:
    while_block_info: Any = ...
    def __init__(
        self,
        auto_parallel_main_prog: Any,
        auto_parallel_startup_prog: Any,
        rank_id: Any,
        dist_context: Any,
        dist_params_grads: Any,
        batch_size: Optional[Any] = ...,
    ) -> None: ...
    @property
    def auto_parallel_main_prog(self): ...
    @property
    def auto_parallel_startup_prog(self): ...
    @property
    def rank_id(self): ...
    @property
    def dist_context(self): ...
    @property
    def dist_params_grads(self): ...
    @property
    def batch_size(self): ...
    @property
    def has_sent(self): ...
    @property
    def has_recv(self): ...
    @property
    def has_allgather(self): ...
    @staticmethod
    def compute_partition_shape(complete_shape: Any, dims_mapping: Any, process_shape: Any): ...
    @staticmethod
    def compute_process_index(process: Any, process_group: Any, process_shape: Any): ...
    @staticmethod
    def compute_partition_index(
        process: Any, complete_shape: Any, dims_mapping: Any, process_shape: Any, process_group: Any
    ): ...
    @staticmethod
    def compute_concat_info(partition_index_x: Any, partition_index_y: Any): ...
    @staticmethod
    def compute_complete_shape(slice_shape: Any, process_shape: Any, dims_mapping: Any): ...
    @staticmethod
    def concat_partitions(partition_index_list: Any, partition_index: Any) -> None: ...
    @staticmethod
    def change_while_op_input_and_output(auto_parallel_main_prog: Any, dist_context: Any) -> None: ...
    def is_overlapped(self, shape_x: Any, shape_y: Any): ...
    def is_unshard(self, dims_mapping: Any): ...
    def is_special_op(self, op: Any): ...
    def is_condition_replicative(self, op: Any): ...
    def need_reshard(self, dist_tensor: Any, dist_op: Any, actual_process_mesh: Any, op_input: bool = ...): ...
    def get_process_meshes(self, op: Any): ...
    def get_op_process_meshes(self, op: Any): ...
    def get_while_op_actual_process_mesh(self, op: Any): ...
    def find_op_desc_seq(self, dist_tensor: Any, dist_op: Any, actual_process_mesh: Any): ...
    def parse_op_desc(self, block: Any, op_desc_seq: Any, var_name: Any, reshard_op: Any, actual_process_mesh: Any): ...
    def reshard(self) -> None: ...
