from __future__ import annotations

from typing import Any

class ParallelMode:
    DATA_PARALLEL: int = ...
    TENSOR_PARALLEL: int = ...
    PIPELINE_PARALLEL: int = ...
    SHARDING_PARALLEL: int = ...

class CommunicateTopology:
    coordinate: Any = ...
    def __init__(self, hybrid_group_names: Any = ..., dims: Any = ...): ...
    def get_hybrid_group_names(self): ...
    def get_dim(self, axis_name: Any): ...
    def world_size(self): ...
    def get_rank(self, **args: Any): ...
    def get_coord(self, rank: Any): ...
    def get_axis_list(self, axis_name: Any, index: Any): ...
    def get_dim_size(self, axis_name: Any): ...
    def get_comm_list(self, axis_name: Any): ...
    def get_rank_from_stage(self, global_rank: Any, **kwargs: Any): ...

class HybridCommunicateGroup:
    nranks: Any = ...
    global_rank: Any = ...
    stage_id: Any = ...
    is_first_stage: Any = ...
    is_last_stage: Any = ...
    def __init__(self, topology: Any) -> None: ...
    def get_parallel_mode(self): ...
    def topology(self): ...
    def get_global_rank(self): ...
    def get_data_parallel_rank(self): ...
    def get_data_parallel_world_size(self): ...
    def get_data_parallel_group(self): ...
    def get_data_parallel_group_src_rank(self): ...
    def get_model_parallel_rank(self): ...
    def get_model_parallel_world_size(self): ...
    def get_model_parallel_group(self): ...
    def get_model_parallel_group_src_rank(self): ...
    def get_stage_id(self): ...
    def get_pipe_parallel_world_size(self): ...
    def get_pipe_parallel_group(self): ...
    def get_p2p_groups(self): ...
    def get_sharding_parallel_rank(self): ...
    def get_sharding_parallel_world_size(self): ...
    def get_sharding_parallel_group(self): ...
    def get_sharding_parallel_group_src_rank(self): ...
    def get_check_parallel_group(self): ...
    def get_rank_from_stage(self, stage_id: Any, **kwargs: Any): ...

class _CommunicateGroup:
    groups: Any = ...
    def __init__(self) -> None: ...
    def set_comm_group(
        self, group_name: Any, group_rank: Any, group_size: Any, ring_id: Any, group_ranks: Any
    ) -> None: ...
    def get_group(self, group_name: Any): ...
    def get_model_parallel_group(self): ...
    def get_model_parallel_world_size(self): ...
    def get_model_parallel_rank(self): ...
