LocalTensorLayout#

class cuquantum.tensornet.experimental.distributed.LocalTensorLayout( rank: int, process_coords: tuple[int, ...], global_shape: tuple[int, ...], local_shape: tuple[int, ...], element_strides: tuple[int, ...], storage_size_elements: int, _process_grid_shape: tuple[int, ...], _block_sizes: tuple[int | None, ...], _first_process: tuple[int, ...], )[source]#

Compact rank-local layout for a block-cyclic distribution.

Methods

__init__( rank: int, process_coords: tuple[int, ...], global_shape: tuple[int, ...], local_shape: tuple[int, ...], element_strides: tuple[int, ...], storage_size_elements: int, _process_grid_shape: tuple[int, ...], _block_sizes: tuple[int | None, ...], _first_process: tuple[int, ...], ) → None#

global_to_local( global_coord: Sequence[int], ) → tuple[int, ...] | None[source]#: Map a global coordinate to this rank’s local coordinate, if owned.

local_to_global( local_coord: Sequence[int], ) → tuple[int, ...][source]#: Map a local coordinate owned by this rank to a global coordinate.

physical_offset( local_coord: Sequence[int], ) → int[source]#: Return the dense element offset of a local coordinate.

Attributes

owned_global_segments#

Per-mode owned global index segments as (start, extent) pairs.

Element ownership is the Cartesian product of these mode-wise segments. Cyclic modes may contribute multiple disjoint segments; slab and replicated modes contribute at most one.

rank: int#

process_coords: tuple[int, ...]#

global_shape: tuple[int, ...]#

local_shape: tuple[int, ...]#

element_strides: tuple[int, ...]#

storage_size_elements: int#