Function distributed_array

cupyx/distributed/array/_array.py:837–915

Creates a distributed array from the given data. This function does not check if all elements of the given array are stored in some of the chunks. Args: array (array_like): :class:`DistributedArray` object, :class:`cupy.ndarray` object or any other object that can b

(
    array: ArrayLike,
    index_map: dict[int, Any],
    mode: _modes.Mode = _modes.REPLICA,
)

def distributed_array(
    array: ArrayLike,
    index_map: dict[int, Any],
    mode: _modes.Mode = _modes.REPLICA,
) -> DistributedArray:
    """Creates a distributed array from the given data.

    This function does not check if all elements of the given array are stored
    in some of the chunks.

    Args:
        array (array_like): :class:`DistributedArray` object,
            :class:`cupy.ndarray` object or any other object that can be passed
            to :func:`numpy.array`.
        index_map (dict from int to array indices): Indices for the chunks
            that devices with designated IDs own. One device can have multiple
            chunks, which can be specified as a list of array indices.
        mode (mode object, optional): Mode that determines how overlaps
            of the chunks are interpreted. Defaults to
            ``cupyx.distributed.array.REPLICA``.

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

    Example:
        >>> array = cupy.arange(9).reshape(3, 3)
        >>> A = distributed_array(
        ...     array,
        ...     {0: [(slice(2), slice(2)),  # array[:2, :2]
        ...          slice(None, None, 2)], # array[::2]
        ...      1:  (slice(1, None), 2)})  # array[1:, 2]
    """
    if isinstance(array, DistributedArray):
        if array.mode != mode:
            array = array.change_mode(mode)
        if array.index_map != index_map:
            array = array.reshard(index_map)
        return DistributedArray(
            array.shape, array.dtype, array._chunks_map, array._mode,
            array._comms)

    if isinstance(array, (numpy.ndarray, ndarray)):
        if mode != _modes.REPLICA:
            array = array.copy()
    else:
        array = numpy.array(array)

    index_map = _index_arith._normalize_index_map(array.shape, index_map)
    comms = None

    # Define how to form a chunk from (dev, idx, src_array)
    make_chunk: Callable[
        [int, int, tuple[slice, ...], ndarray, list[Any] | None],
        _Chunk
    ]

    if isinstance(array, ndarray):
        src_dev = array.device.id