def test_dist_sizes(self): dist = Distribution(self.context, (2, 3, 4), dist=('n', 'b', 'c')) ddpr = dist.get_dim_data_per_rank() shapes = metadata_utils.shapes_from_dim_data_per_rank(ddpr) if self.context.nengines == 4: self.assertEqual(shapes, [(2, 2, 2), (2, 2, 2), (2, 1, 2), (2, 1, 2)])
def test_c_size(self): dim_dict = {'dist_type': 'c', 'size': 42, 'proc_grid_size': 2, 'proc_grid_rank': 0, 'start': 0} dist = Distribution.from_global_dim_data(self.context, (dim_dict,)) ddpr = dist.get_dim_data_per_rank() shapes = metadata_utils.shapes_from_dim_data_per_rank(ddpr) self.assertEqual(shapes, [(21,), (21,)])
def localshapes(self): return shapes_from_dim_data_per_rank(self.get_dim_data_per_rank())