def test_block_redist_2D_one_to_many(self):
source_dist = Distribution(self.context, (9, 9), ('b', 'b'), (1, 1),
targets=[2])
dest_dist = Distribution(self.context, (9, 9), ('b', 'b'), (2, 2),
targets=range(4))
plan = source_dist.get_redist_plan(dest_dist)
expected = [
{
'source_rank': 2,
'dest_rank': 0,
'indices': [(0, 5, 1), (0, 5, 1)]
},
{
'source_rank': 2,
'dest_rank': 1,
'indices': [(0, 5, 1), (5, 9, 1)]
},
{
'source_rank': 2,
'dest_rank': 2,
'indices': [(5, 9, 1), (0, 5, 1)]
},
{
'source_rank': 2,
'dest_rank': 3,
'indices': [(5, 9, 1), (5, 9, 1)]
},
]
for p, e in zip(plan, expected):
self.assertEqual(p, e)
def test_block_redist_many_to_one(self):
source_dist = Distribution(self.context, (40,), ("b",), (2,), targets=[1, 2])
dest_dist = Distribution(self.context, (40,), ("b",), (1,), targets=[0])
plan = source_dist.get_redist_plan(dest_dist)
expected = [
{"source_rank": 1, "dest_rank": 0, "indices": [(0, 20, 1)]},
{"source_rank": 2, "dest_rank": 0, "indices": [(20, 40, 1)]},
]
self.assertEqual(plan, expected)
def test_block_redist_2D_one_to_many(self):
source_dist = Distribution(self.context, (9, 9), ("b", "b"), (1, 1), targets=[2])
dest_dist = Distribution(self.context, (9, 9), ("b", "b"), (2, 2), targets=range(4))
plan = source_dist.get_redist_plan(dest_dist)
expected = [
{"source_rank": 2, "dest_rank": 0, "indices": [(0, 5, 1), (0, 5, 1)]},
{"source_rank": 2, "dest_rank": 1, "indices": [(0, 5, 1), (5, 9, 1)]},
{"source_rank": 2, "dest_rank": 2, "indices": [(5, 9, 1), (0, 5, 1)]},
{"source_rank": 2, "dest_rank": 3, "indices": [(5, 9, 1), (5, 9, 1)]},
]
for p, e in zip(plan, expected):
self.assertEqual(p, e)
def test_block_redist_many_to_one(self):
source_dist = Distribution(self.context, (40, ), ('b', ), (2, ),
targets=[1, 2])
dest_dist = Distribution(self.context, (40, ), ('b', ), (1, ),
targets=[0])
plan = source_dist.get_redist_plan(dest_dist)
expected = [
{
'source_rank': 1,
'dest_rank': 0,
'indices': [(0, 20, 1)]
},
{
'source_rank': 2,
'dest_rank': 0,
'indices': [(20, 40, 1)]
},
]
self.assertEqual(plan, expected)
