from __future__ import annotations

from datetime import timedelta as timedelta
from typing import Any, Optional

from ..fluid.data_feeder import check_dtype as check_dtype
from ..fluid.data_feeder import check_type as check_type
from ..fluid.data_feeder import check_variable_and_dtype as check_variable_and_dtype
from ..fluid.data_feeder import convert_dtype as convert_dtype
from ..fluid.dygraph import layers as layers
from ..fluid.dygraph.parallel import prepare_context as prepare_context
from ..fluid.framework import OpProtoHolder as OpProtoHolder
from ..fluid.framework import Variable as Variable
from ..fluid.framework import convert_np_dtype_to_dtype_ as convert_np_dtype_to_dtype_
from ..fluid.framework import in_dygraph_mode as in_dygraph_mode
from ..fluid.layer_helper import LayerHelper as LayerHelper
from ..fluid.layers import utils as utils
from ..fluid.layers.tensor import fill_constant as fill_constant

class ReduceOp:
    SUM: int = ...
    MAX: int = ...
    MIN: int = ...
    PROD: int = ...
    AVG: int = ...

class Group:
    rank: Any = ...
    nranks: Any = ...
    id: Any = ...
    ranks: Any = ...
    pg: Any = ...
    name: Any = ...
    def __init__(
        self,
        rank: Any,
        rank_num: Any,
        id: int = ...,
        ranks: Any = ...,
        pg: Optional[Any] = ...,
        name: Optional[Any] = ...,
    ) -> None: ...
    def is_member(self): ...
    def get_group_rank(self, rank: Any): ...
    @property
    def process_group(self): ...

def get_group(id: int = ...): ...
def barrier(group: Optional[Any] = ...): ...
def new_group(ranks: Optional[Any] = ..., backend: Optional[Any] = ...): ...
def wait(tensor: Any, group: Optional[Any] = ..., use_calc_stream: bool = ...) -> None: ...
def broadcast(tensor: Any, src: Any, group: Optional[Any] = ..., use_calc_stream: bool = ...): ...
def all_reduce(tensor: Any, op: Any = ..., group: Optional[Any] = ..., use_calc_stream: bool = ...): ...
def reduce(tensor: Any, dst: Any, op: Any = ..., group: Optional[Any] = ..., use_calc_stream: bool = ...): ...
def all_gather(tensor_list: Any, tensor: Any, group: Optional[Any] = ..., use_calc_stream: bool = ...) -> None: ...
def scatter(
    tensor: Any,
    tensor_list: Optional[Any] = ...,
    src: int = ...,
    group: Optional[Any] = ...,
    use_calc_stream: bool = ...,
): ...

class _Linear(layers.Layer):
    weight: Any = ...
    bias: Any = ...
    name: Any = ...
    def __init__(
        self,
        in_features: Any,
        out_features: Any,
        weight_attr: Optional[Any] = ...,
        bias_attr: Optional[Any] = ...,
        name: Optional[Any] = ...,
    ) -> None: ...
    def forward(self, input: Any): ...
    def extra_repr(self): ...

def split(
    x: Any,
    size: Any,
    operation: Any,
    axis: int = ...,
    num_partitions: int = ...,
    gather_out: bool = ...,
    weight_attr: Optional[Any] = ...,
    bias_attr: Optional[Any] = ...,
    name: Optional[Any] = ...,
): ...
def alltoall(
    in_tensor_list: Any, out_tensor_list: Any, group: Optional[Any] = ..., use_calc_stream: bool = ...
) -> None: ...
def send(tensor: Any, dst: int = ..., group: Optional[Any] = ..., use_calc_stream: bool = ...): ...
def recv(tensor: Any, src: int = ..., group: Optional[Any] = ..., use_calc_stream: bool = ...): ...
