qimpy.mpi.ProcessGrid

class ProcessGrid(dim_names, shape=None)

Bases: object

Process grid of shape dimensions over the global process group. Any -1 entries in shape are undetermined and will be resolved after the number of tasks split along that dimension are set using provide_n_tasks. Subsequently, use get_group to get arbitrary hyperplane process groups that connect processes whose index only varies along specified subsets of dimensions.

Parameters:

dim_names (str)
shape (ndarray)

__init__(dim_names, shape=None)

Parameters:

dim_names (str)
shape (Sequence[int] | None)

Return type:

None

Methods

`__init__`
`get_group`	Get communicator for a hyper-plane spanning dim_names.
`get_submesh`	Get device mesh with a subset of dimensions.
`provide_n_tasks`	Provide task count for a process grid dimension named dim_name.

Attributes

`dim_names`	Each character (must be unique) names a dimension.
`shape`	Grid dimensions.
`is_resolved`	Whether all dimensions have been resolved
`device_mesh`	Corresponding device mesh

get_group(dim_names)

Get communicator for a hyper-plane spanning dim_names. The resulting communicator will connect processes whose index in the process grid only varies along dimensions within dim_names. All dimensions should be fully resolved before getting any groups.

Parameters:: dim_names (str)
Return type:: ProcessGroup

get_submesh(dim_names)

Get device mesh with a subset of dimensions.

Parameters:: dim_names (str)
Return type:: DeviceMesh

provide_n_tasks(dim_name, n_tasks)

Provide task count for a process grid dimension named dim_name. If that dimension is undetermined (-1), set it to a suitable value that is compatible with the total processes, any other known dimensions, and with splitting n_tasks tasks with reasonable load balancing over this dimension.

Parameters:

dim_name (str) – Name of dimension (single charcater) to provide n_tasks for.
n_tasks (int) – Number of tasks available to split on this dimension of the process grid, used for setting dimension to ensure reasonable load balancing.

Return type:

None

device_mesh: DeviceMesh: Corresponding device mesh

dim_names: str: Each character (must be unique) names a dimension.

is_resolved: bool: Whether all dimensions have been resolved

shape: ndarray: Grid dimensions. Unresolved dimensions are -1.