from __future__ import annotations

from typing import Any, Optional

from ...utils.log_util import logger as logger
from .utils import number_2_dtype as number_2_dtype
from .utils import paddle_2_number as paddle_2_number

def initialize_p2p_groups(hcg: Any, use_cache: bool = ...) -> None: ...

class SendRecvMeta:
    send_shape_message: Any = ...
    send_dtype_message: Any = ...
    recv_shape_message: Any = ...
    recv_dtype_message: Any = ...
    recv_stop_gradient: Any = ...
    has_send_meta: bool = ...
    has_recv_meta: bool = ...
    def __init__(self) -> None: ...
    def recv_meta(self, group: Any) -> None: ...
    def send_meta(self, tensor: Any, group: Any) -> None: ...
    def set_send_message(self, tensor: Any) -> None: ...

def send_partial(
    tensor: Any,
    dst: int = ...,
    nranks: int = ...,
    rank_id: int = ...,
    group: Optional[Any] = ...,
    use_calc_stream: bool = ...,
): ...
def recv_partial(
    tensor: Any,
    src: int = ...,
    nranks: int = ...,
    rank_id: int = ...,
    group: Optional[Any] = ...,
    use_calc_stream: bool = ...,
) -> None: ...
def allgather_partial(
    tensor: Any, nranks: int = ..., rank_id: int = ..., group: Optional[Any] = ..., use_calc_stream: bool = ...
): ...
def recv_forward(): ...
def recv_backward(): ...
def send_forward(output_tensor: Any) -> None: ...
def send_backward(input_tensor_grad: Any) -> None: ...
def send_forward_recv_backward(output_tensor: Any): ...
def send_backward_recv_forward(input_tensor_grad: Any): ...
