cupyx.distributed.array.DistributedArray

class cupyx.distributed.array.DistributedArray(self, shape, dtype, chunks_map, mode=REPLICA, comms=None)

Multi-dimensional array distributed across multiple CUDA devices.

This class implements some elementary operations that 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 distributed_array().

Parameters:

• 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.

Return type:

DistributedArray

See also

DistributedArray.mode for details about modes.

Methods

__getitem__(*args, **kwargs)

Not supported.

__setitem__(*args, **kwargs)

Not supported.

__len__(*args, **kwargs)

Not supported.

__iter__(*args, **kwargs)

Not supported.

__copy__(*args, **kwargs)

Not supported.

all(*args, **kwargs)

Not supported.

all_chunks()

Return the chunks with all buffered data flushed.

Buffered data are created in situations such as resharding and mode changing.

Return type:

dict[int, list[cupy.ndarray]]

any(*args, **kwargs)

Not supported.

argmax(*args, **kwargs)

Not supported.

argmin(*args, **kwargs)

Not supported.

argpartition(*args, **kwargs)

Not supported.

argsort(*args, **kwargs)

Not supported.

astype(*args, **kwargs)

Not supported.

change_mode(mode)

Return a view or a copy in the given mode.

Parameters:

mode (mode Object) – How overlaps of the chunks are interpreted.

Return type:

DistributedArray

See also

DistributedArray.mode for details about modes.

choose(*args, **kwargs)

Not supported.

clip(*args, **kwargs)

Not supported.

compress(*args, **kwargs)

Not supported.

conj(self) → ndarray

conjugate(self) → ndarray

copy(*args, **kwargs)

Not supported.

cumprod(*args, **kwargs)

Not supported.

cumsum(*args, **kwargs)

Not supported.

diagonal(*args, **kwargs)

Not supported.

dot(*args, **kwargs)

Not supported.

dump(*args, **kwargs)

Not supported.

dumps(*args, **kwargs)

Not supported.

fill(*args, **kwargs)

Not supported.

flatten(*args, **kwargs)

Not supported.

get(stream=None, order='C', out=None, blocking=True)

Return a copy of the array on the host memory.

Return type:

ndarray

item(*args, **kwargs)

Not supported.

max(axis=None, out=None, keepdims=False)

Return the maximum along a given axis.

Note

Currently, it only supports non-None values for axis and the default values for out and keepdims.

See also

cupy.ndarray.max(), numpy.ndarray.max()

mdspan(self, *, index_type, allow_unsafe=False)

Returns an mdspan view of the array for use in CUDA kernels.

This method creates a view of the CuPy array that is compatible with cuda::std::mdspan for use in custom CUDA kernels.

Parameters:

• index_type (dtype) – The data type for extent and stride indices. Must be either cupy.int32 or cupy.int64. If cupy.int32 is specified, the array size must not exceed INT32_MAX.

• allow_unsafe (bool) – If True, allows creating an mdspan for arrays that have either zero or negative strides, or one or more dimensions of size zero. Depending on the access pattern such mdspan may lead to undefined behavior when described with either a left-, right, or stride- layout_stride as per C++ standard. For example, mdspan.required_span_size() might become negative. Default is False.

Returns:

An mdspan view of the array that can be passed to CUDA kernels as a kernel argument.

Return type:

mdspan

Raises:

ValueError – If index_type is not cupy.int32 or cupy.int64, or if the array size exceeds the range of the specified index_type.

Note

The returned mdspan can work with either layout_stride, layout_left, or layout_right, but your kernel must declare the mdspan type with all extents being dynamic:

template<typename T, typename IndexType>
__global__ void my_kernel(
    cuda::std::mdspan<
        T,
        cuda::std::extents<
            IndexType,
            cuda::std::dynamic_extent,
            cuda::std::dynamic_extent>,
        cuda::std::layout_stride> arr
) {
    // arr is a 2D mdspan
    // Access: arr(i, j)
}

Static extents are NOT supported. For example, using extents<int, 4, 8> will result in undefined behavior.

Example

>>> import cupy
>>> a = cupy.arange(12, dtype=cupy.float32).reshape(3, 4)
>>> a_mdspan = a.mdspan(index_type=cupy.int64)
>>> # Pass a_mdspan to a cupy.RawKernel expecting:
>>> # mdspan<float, extents<int64_t, dyn, dyn>, layout_stride>

mean(*args, **kwargs)

Not supported.

min(axis=None, out=None, keepdims=False)

Return the minimum along a given axis.

Note

Currently, it only supports non-None values for axis and the default values for out and keepdims.

See also

cupy.ndarray.min(), numpy.ndarray.min()

nonzero(*args, **kwargs)

Not supported.

partition(*args, **kwargs)

Not supported.

prod(axis=None, dtype=None, out=None, keepdims=None)

Return the minimum along a given axis.

Note

Currently, it only supports non-None values for axis and the default values for out and keepdims.

See also

cupy.ndarray.prod(), numpy.ndarray.prod()

ptp(*args, **kwargs)

Not supported.

put(*args, **kwargs)

Not supported.

ravel(*args, **kwargs)

Not supported.

reduced_view(*args, **kwargs)

Not supported.

repeat(*args, **kwargs)

Not supported.

reshape(*args, **kwargs)

Not supported.

reshard(index_map)

Return a view or a copy having the given index_map.

Data transfers across devices are done on separate streams created internally. To make them asynchronous, transferred data is buffered and reflected to the chunks when necessary.

Parameters:

index_map (dict from int to array indices) – Indices for the chunks that devices with designated IDs own. The current index_map of a distributed array can be obtained from DistributedArray.index_map.

Return type:

DistributedArray

round(*args, **kwargs)

Not supported.

scatter_add(*args, **kwargs)

Not supported.

scatter_max(*args, **kwargs)

Not supported.

scatter_min(*args, **kwargs)

Not supported.

searchsorted(*args, **kwargs)

Not supported.

set(*args, **kwargs)

Not supported.

sort(*args, **kwargs)

Not supported.

squeeze(*args, **kwargs)

Not supported.

std(*args, **kwargs)

Not supported.

sum(axis=None, dtype=None, out=None, keepdims=False)

Return the minimum along a given axis.

Note

Currently, it only supports non-None values for axis and the default values for out and keepdims.

See also

cupy.ndarray.sum(), numpy.ndarray.sum()

swapaxes(*args, **kwargs)

Not supported.

take(*args, **kwargs)

Not supported.

toDlpack(*args, **kwargs)

Not supported.

tobytes(*args, **kwargs)

Not supported.

tofile(*args, **kwargs)

Not supported.

tolist(*args, **kwargs)

Not supported.

trace(*args, **kwargs)

Not supported.

transpose(*args, **kwargs)

Not supported.

var(*args, **kwargs)

Not supported.

view(*args, **kwargs)

Not supported.

__eq__(value, /)

Return self==value.

__ne__(value, /)

Return self!=value.

__lt__(value, /)

Return self<value.

__le__(value, /)

Return self<=value.

__gt__(value, /)

Return self>value.

__ge__(value, /)

Return self>=value.

__bool__()

True if self else False

Attributes

T

Not supported.

base

Not supported.

cstruct

Not supported.

data

Not supported.

device

Not supported.

devices

A collection of device IDs holding part of the data.

dtype

flags

Not supported.

flat

Not supported.

imag

Not supported.

index_map

Indices for the chunks that devices with designated IDs own.

itemsize

Size of each element in bytes.

See also

numpy.ndarray.itemsize

mT

Matrix-transpose view of the array.

If ndim < 2, raise a ValueError.

mode

Describe how overlaps of the chunks are interpreted.

In the replica mode, chunks are guaranteed to have identical values on their overlapping segments. In other modes, they are not necessarily identical and represent the original data as their max, sum, etc.

DistributedArray currently supports cupyx.distributed.array.REPLICA, cupyx.distributed.array.MIN, cupyx.distributed.array.MAX, cupyx.distributed.array.SUM, cupyx.distributed.array.PROD modes.

Many operations on distributed arrays including cupy.ufunc and matmul() involve changing their mode beforehand. These mode conversions are done automatically, so in most cases users do not have to manage modes manually.

Example

>>> A = distributed_array(
...     cupy.arange(6).reshape(2, 3),
...     make_2d_index_map([0, 2], [0, 1, 3],
...                       [[{0}, {1, 2}]]))
>>> B = distributed_array(
...     cupy.arange(12).reshape(3, 4),
...     make_2d_index_map([0, 1, 3], [0, 2, 4],
...                       [[{0}, {0}],
...                        [{1}, {2}]]))
>>> C = A @ B
>>> C
array([[20, 23, 26, 29],
       [56, 68, 80, 92]])
>>> C.mode
'sum'
>>> C.all_chunks()
{0: [array([[0, 0],
            [0, 3]]),     # left half
     array([[0, 0],
            [6, 9]])],    # right half
 1: [array([[20, 23],
            [56, 65]])],  # left half
 2: [array([[26, 29],
            [74, 83]])]}  # right half
>>> C_replica = C.change_mode('replica')
>>> C_replica.mode
'replica'
>>> C_replica.all_chunks()
{0: [array([[20, 23],
            [56, 68]]),   # left half
     array([[26, 29],
            [80, 92]])],  # right half
 1: [array([[20, 23],
            [56, 68]])],  # left half
 2: [array([[26, 29],
            [80, 92]])]}  # right half

nbytes

Total size of all elements in bytes.

It does not count skips between elements.

See also

numpy.ndarray.nbytes

ndim

Number of dimensions.

a.ndim is equivalent to len(a.shape).

See also

numpy.ndarray.ndim

real

Not supported.

shape

Tuple of array dimensions.

Assignment to this property is currently not supported.

size

strides

Not supported.