import torch
import numpy as np
from mpi4py import MPI

from parallel_pytorch.ops import tensor_merge
from parallel_pytorch.utils import abort_on_exception


@abort_on_exception
def test_1():
    worker_shape = [2, 2]

    world = MPI.COMM_WORLD
    num_workers = np.array(worker_shape).prod()
    comm = MPI.COMM_WORLD.Split(color=0 if world.Get_rank() < num_workers else 1, key=world.Get_rank())
    if world.Get_rank() < num_workers:
        if comm.Get_rank() == 0:
            x = torch.tensor([[0, 1], [4, 5]])
        elif comm.Get_rank() == 1:
            x = torch.tensor([[2, 3], [6, 7]])
        elif comm.Get_rank() == 2:
            x = torch.tensor([[8, 9], [12, 13]])
        elif comm.Get_rank() == 3:
            x = torch.tensor([[10, 11], [14, 15]])
        x = tensor_merge(x, comm=comm, worker_shape=worker_shape)
        if comm.Get_rank() == 0:
            e = torch.tensor([
                [0, 1, 2, 3],
                [4, 5, 6, 7],
                [8, 9, 10, 11],
                [12, 13, 14, 15],
            ])
            assert torch.allclose(x, e), f'{x} != {e}'


@abort_on_exception
def test_2():
    x_shape = [2, 2]
    worker_shape = [1, 1]

    world = MPI.COMM_WORLD
    num_workers = np.array(worker_shape).prod()
    comm = MPI.COMM_WORLD.Split(color=0 if world.Get_rank() < num_workers else 1, key=world.Get_rank())
    if world.Get_rank() < num_workers:
        volume = np.array(x_shape).prod()
        x = torch.arange(volume).view(x_shape)
        x = tensor_merge(x, comm=comm, worker_shape=worker_shape)
        e = torch.tensor([[0, 1], [2, 3]])
        assert torch.allclose(x, e), f'{x} != {e}'


def run_all():
    test_1()
    test_2()


if __name__ == '__main__':
    run_all()

[((25, 12, 25, 65), 'parallel_pytorch.ops.tensor_merge', 'tensor_merge', (), '', False, 'from parallel_pytorch.ops import tensor_merge\n'), ((47, 12, 47, 65), 'parallel_pytorch.ops.tensor_merge', 'tensor_merge', (), '', False, 'from parallel_pytorch.ops import tensor_merge\n'), ((48, 12, 48, 42), 'torch.tensor', 'torch.tensor', ({(48, 25, 48, 41): '[[0, 1], [2, 3]]'}, {}), '([[0, 1], [2, 3]])', False, 'import torch\n'), ((49, 15, 49, 35), 'torch.allclose', 'torch.allclose', ({(49, 30, 49, 31): 'x', (49, 33, 49, 34): 'e'}, {}), '(x, e)', False, 'import torch\n'), ((14, 18, 14, 40), 'numpy.array', 'np.array', ({(14, 27, 14, 39): 'worker_shape'}, {}), '(worker_shape)', True, 'import numpy as np\n'), ((18, 16, 18, 46), 'torch.tensor', 'torch.tensor', ({(18, 29, 18, 45): '[[0, 1], [4, 5]]'}, {}), '([[0, 1], [4, 5]])', False, 'import torch\n'), ((27, 16, 32, 14), 'torch.tensor', 'torch.tensor', ({(27, 29, 32, 13): '[[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15]]'}, {}), '([[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15]])', False, 'import torch\n'), ((33, 19, 33, 39), 'torch.allclose', 'torch.allclose', ({(33, 34, 33, 35): 'x', (33, 37, 33, 38): 'e'}, {}), '(x, e)', False, 'import torch\n'), ((42, 18, 42, 40), 'numpy.array', 'np.array', ({(42, 27, 42, 39): 'worker_shape'}, {}), '(worker_shape)', True, 'import numpy as np\n'), ((20, 16, 20, 46), 'torch.tensor', 'torch.tensor', ({(20, 29, 20, 45): '[[2, 3], [6, 7]]'}, {}), '([[2, 3], [6, 7]])', False, 'import torch\n'), ((45, 17, 45, 34), 'numpy.array', 'np.array', ({(45, 26, 45, 33): 'x_shape'}, {}), '(x_shape)', True, 'import numpy as np\n'), ((46, 12, 46, 32), 'torch.arange', 'torch.arange', ({(46, 25, 46, 31): 'volume'}, {}), '(volume)', False, 'import torch\n'), ((22, 16, 22, 48), 'torch.tensor', 'torch.tensor', ({(22, 29, 22, 47): '[[8, 9], [12, 13]]'}, {}), '([[8, 9], [12, 13]])', False, 'import torch\n'), ((24, 16, 24, 50), 'torch.tensor', 'torch.tensor', ({(24, 29, 24, 49): '[[10, 11], [14, 15]]'}, {}), '([[10, 11], [14, 15]])', False, 'import torch\n')]