import torch
from typing import Any, Dict, Optional, Union

from .parallel_state import get_tp_group

def tensor_model_parallel_all_reduce(input_: torch.Tensor, tp_group: Any = None) -> torch.Tensor:
    """All-reduce the input tensor across model parallel group."""
    return get_tp_group(tp_group).all_reduce(input_)


def tensor_model_parallel_all_gather(input_: torch.Tensor,
                                     dim: int = -1, tp_group: Any = None) -> torch.Tensor:
    """All-gather the input tensor across model parallel group."""
    return get_tp_group(tp_group).all_gather(input_, dim)


def tensor_model_parallel_gather(input_: torch.Tensor,
                                 dst: int = 0,
                                 dim: int = -1, tp_group: Any = None) -> Optional[torch.Tensor]:
    """Gather the input tensor across model parallel group."""
    return get_tp_group(tp_group).gather(input_, dst, dim)
