BlockCyclic#

class cuquantum.tensornet.experimental.distributed.BlockCyclic( process_grid: ProcessGrid, block_sizes: Sequence[int | None], *, first_process: Sequence[int] | None = None, )[source]#

N-D block-cyclic distribution with per-mode slab support.

Each tensor mode is partitioned across the corresponding process-grid dimension. An integer block size assigns consecutive blocks round-robin across grid coordinates (the final block may be shorter), while None assigns one near-even contiguous slab per coordinate. A grid extent of one leaves that mode undistributed.

Parameters:

process_grid – The ProcessGrid of this subpackage (the nvmath parent class is not accepted).
block_sizes – Per-mode entries, one per process-grid dimension: a positive integer for block-cyclic partitioning, or None for slab partitioning.
first_process – Optional grid coordinates of the process owning the first block of each mode; defaults to the grid origin.

Warning

This API is experimental and subject to future changes.

Methods

__init__( process_grid: ProcessGrid, block_sizes: Sequence[int | None], *, first_process: Sequence[int] | None = None, )[source]#

copy() → Distribution#: This is a common implementation for those distributions that only require a shallow copy.

shape( rank: int, global_shape: Sequence[int] | None = None, ) → tuple[int, ...][source]#

Return the local shard shape for rank and global_shape.

When global_shape is omitted, use the bound global shape.

to(cls, /, *, ndim=None, copy=False)[source]#

Convert to a compatible nvmath distribution type.

Block-cyclic conversions return this distribution (or a copy). Conversion to Slab requires exactly one unshifted slab-distributed mode and no replica layers. Other conversions are not supported.

Attributes

block_sizes#: Per-mode block size; None denotes contiguous slab partitioning.

first_process#: Index in the process grid of the process who owns the first block of the distributed tensor.

ndim#

process_grid#: The process grid of this BlockCyclic distribution