Class DistributedArray

cupyx/distributed/array/_array.py:65–834

__init__(self, shape, dtype, chunks_map, mode=REPLICA, comms=None) Multi-dimensional array distributed across multiple CUDA devices. This class implements some elementary operations that :class:`cupy.ndarray` provides. The array content is split into chunks, contiguous arrays

class DistributedArray(ndarray):
    """
    __init__(self, shape, dtype, chunks_map, mode=REPLICA, comms=None)

    Multi-dimensional array distributed across multiple CUDA devices.

    This class implements some elementary operations that :class:`cupy.ndarray`
    provides. The array content is split into chunks, contiguous arrays
    corresponding to slices of the original array. Note that one device can
    hold multiple chunks.

    This direct constructor is designed for internal calls. Users should create
    distributed arrays using :func:`distributed_array`.

    Args:
        shape (tuple of ints): Shape of created array.
        dtype (dtype_like): Any object that can be interpreted as a numpy data
            type.
        chunks_map (dict from int to list of chunks): Lists of chunk objects
            associated with each device.
        mode (mode object, optional): Mode that determines how overlaps
            of the chunks are interpreted. Defaults to
            ``cupyx.distributed.array.REPLICA``.
        comms (optional): Communicator objects which a distributed array
            hold internally. Sharing them with other distributed arrays can
            save time because their initialization is a costly operation.

    .. seealso::
            :attr:`DistributedArray.mode` for details about modes.
    """

    _chunks_map: dict[int, list[_Chunk]]
    _mode: _modes.Mode
    _streams: dict[int, Stream]
    _comms: dict[int, _Communicator]

    def __new__(  # noqa: PYI034
        cls, shape: tuple[int, ...], dtype: DTypeLike,
        chunks_map: dict[int, list[_Chunk]],
        mode: _modes.Mode = _modes.REPLICA,
        comms: dict[int, _Communicator] | None = None,
    ) -> DistributedArray:
        mem = _MultiDeviceDummyMemory(0)
        memptr = _MultiDeviceDummyPointer(mem, 0)
        obj = super().__new__(cls, shape, dtype, memptr=memptr)
        obj._chunks_map = chunks_map

        obj._mode = mode

        obj._streams = {}
        obj._comms = comms if comms is not None else {}

        return obj

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)

    def __array_finalize__(self, obj):