def from_shape(cls, comm, shape, dist=None, grid_shape=None):
"""Create a Distribution from a `shape` and optional arguments."""
dist = {0: 'b'} if dist is None else dist
ndim = len(shape)
dist_tuple = normalize_dist(dist, ndim)
base_comm = construct.init_base_comm(comm)
comm_size = base_comm.Get_size()
if grid_shape is None: # Make a new grid_shape if not provided.
grid_shape = make_grid_shape(shape, dist_tuple, comm_size)
grid_shape = normalize_grid_shape(grid_shape, shape,
dist_tuple, comm_size)
comm = construct.init_comm(base_comm, grid_shape)
grid_coords = comm.Get_coords(comm.Get_rank())
dim_data = []
for dist, size, grid_rank, grid_size in zip(dist_tuple, shape,
grid_coords, grid_shape):
dim_dict = dict(dist_type=dist,
size=size,
proc_grid_rank=grid_rank,
proc_grid_size=grid_size)
distribute_indices(dim_dict)
dim_data.append(dim_dict)
return cls(comm=base_comm, dim_data=dim_data)
def from_shape(cls, comm, shape, dist=None, grid_shape=None):
"""Create a Distribution from a `shape` and optional arguments."""
dist = {0: 'b'} if dist is None else dist
ndim = len(shape)
dist_tuple = normalize_dist(dist, ndim)
base_comm = construct.init_base_comm(comm)
comm_size = base_comm.Get_size()
if grid_shape is None: # Make a new grid_shape if not provided.
grid_shape = make_grid_shape(shape, dist_tuple, comm_size)
grid_shape = normalize_grid_shape(grid_shape, shape, dist_tuple,
comm_size)
comm = construct.init_comm(base_comm, grid_shape)
grid_coords = comm.Get_coords(comm.Get_rank())
dim_data = []
for dist, size, grid_rank, grid_size in zip(dist_tuple, shape,
grid_coords, grid_shape):
dim_dict = dict(dist_type=dist,
size=size,
proc_grid_rank=grid_rank,
proc_grid_size=grid_size)
distribute_indices(dim_dict)
dim_data.append(dim_dict)
return cls(comm=base_comm, dim_data=dim_data)
def __init__(self, comm, dim_data):
"""Create a Distribution from a `dim_data` structure."""
self._maps = tuple(map_from_dim_dict(dim_dict) for dim_dict in dim_data)
self.base_comm = construct.init_base_comm(comm)
self.comm = construct.init_comm(self.base_comm, self.grid_shape)
def __init__(self, comm, dim_data):
"""Create a Distribution from a `dim_data` structure."""
self._maps = tuple(
map_from_dim_dict(dim_dict) for dim_dict in dim_data)
self.base_comm = construct.init_base_comm(comm)
self.comm = construct.init_comm(self.base_comm, self.grid_shape)