ProcessGrid#

class cuquantum.tensornet.experimental.distributed.ProcessGrid( *, shape: Sequence[int] | None = None, layout: Layout | None = None, process_array=None, )[source]#

An N-dimensional regular process grid with complete replica layers.

Ranks fill the grid in layout order (column- or row-major). When the communicator size exceeds the grid size, it must be an integer multiple of it, and the surplus ranks form complete replica layers: rank r occupies the same coordinates as rank r % size, in replica layer r // size. A mode whose grid extent is one is not partitioned.

Parameters:

shape – Per-mode extents of the process grid.
layout – ProcessGrid.Layout.COL_MAJOR or ROW_MAJOR rank ordering; optional when at most one mode has extent greater than one.
process_array – Custom process arrangements are not supported; must be None.

Warning

This API is experimental and subject to future changes.

Methods

__init__( *, shape: Sequence[int] | None = None, layout: Layout | None = None, process_array=None, )[source]#

Create a new ProcessGrid object.

Parameters:

shape – Shape of the process grid.
layout – Layout of the process grid (column-major or row-major). This is optional for 1D grid or when a custom grid is provided.
process_array – optional ndarray specifying custom arrangement of processes.

coords_to_rank( coords: Sequence[int], replica_id: int = 0, ) → int[source]#: Return the rank for active-grid coordinates and a replica layer.

rank_to_coords(rank: int) → tuple[int, ...][source]#: Return a rank’s coordinates within its active-grid replica.

replica_id(rank: int) → int[source]#: Return a rank’s replica-layer index.

Attributes

layout#: Layout of process grid if row-major or column-major, otherwise None.

process_array#

replication_factor#: Number of complete process-grid replica layers.

shape#: Shape of process grid.

size#: Number of coordinates in one active process grid.