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()

[((781, 834), 'parallel_pytorch.ops.tensor_merge', 'tensor_merge', (['x'], {'comm': 'comm', 'worker_shape': 'worker_shape'}), '(x, comm=comm, worker_shape=worker_shape)\n', (793, 834), False, 'from parallel_pytorch.ops import tensor_merge\n'), ((1498, 1551), 'parallel_pytorch.ops.tensor_merge', 'tensor_merge', (['x'], {'comm': 'comm', 'worker_shape': 'worker_shape'}), '(x, comm=comm, worker_shape=worker_shape)\n', (1510, 1551), False, 'from parallel_pytorch.ops import tensor_merge\n'), ((1564, 1594), 'torch.tensor', 'torch.tensor', (['[[0, 1], [2, 3]]'], {}), '([[0, 1], [2, 3]])\n', (1576, 1594), False, 'import torch\n'), ((1610, 1630), 'torch.allclose', 'torch.allclose', (['x', 'e'], {}), '(x, e)\n', (1624, 1630), False, 'import torch\n'), ((264, 286), 'numpy.array', 'np.array', (['worker_shape'], {}), '(worker_shape)\n', (272, 286), True, 'import numpy as np\n'), ((486, 516), 'torch.tensor', 'torch.tensor', (['[[0, 1], [4, 5]]'], {}), '([[0, 1], [4, 5]])\n', (498, 516), False, 'import torch\n'), ((884, 960), 'torch.tensor', 'torch.tensor', (['[[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]])\n', (896, 960), False, 'import torch\n'), ((1059, 1079), 'torch.allclose', 'torch.allclose', (['x', 'e'], {}), '(x, e)\n', (1073, 1079), False, 'import torch\n'), ((1224, 1246), 'numpy.array', 'np.array', (['worker_shape'], {}), '(worker_shape)\n', (1232, 1246), True, 'import numpy as np\n'), ((568, 598), 'torch.tensor', 'torch.tensor', (['[[2, 3], [6, 7]]'], {}), '([[2, 3], [6, 7]])\n', (580, 598), False, 'import torch\n'), ((1414, 1431), 'numpy.array', 'np.array', (['x_shape'], {}), '(x_shape)\n', (1422, 1431), True, 'import numpy as np\n'), ((1451, 1471), 'torch.arange', 'torch.arange', (['volume'], {}), '(volume)\n', (1463, 1471), False, 'import torch\n'), ((650, 682), 'torch.tensor', 'torch.tensor', (['[[8, 9], [12, 13]]'], {}), '([[8, 9], [12, 13]])\n', (662, 682), False, 'import torch\n'), ((734, 768), 'torch.tensor', 'torch.tensor', (['[[10, 11], [14, 15]]'], {}), '([[10, 11], [14, 15]])\n', (746, 768), False, 'import torch\n')]