# python3.7
"""Contains operators for neural networks."""

import torch
import torch.distributed as dist

__all__ = ['all_gather']


def all_gather(tensor):
    """Gathers tensor from all devices and executes averaging."""
    if not dist.is_initialized():
        return tensor

    world_size = dist.get_world_size()
    tensor_list = [torch.ones_like(tensor) for _ in range(world_size)]
    dist.all_gather(tensor_list, tensor, async_op=False)
    return torch.stack(tensor_list, dim=0).mean(dim=0)