cupyx.distributed.array.matmul#

cupyx.distributed.array.matmul(a, b, out=None, **kwargs)[source]#

Matrix multiplication between distributed arrays.

The arguments must have compatible shape and index_map.

This operation converts its operands into the replica mode, and compute their product in the sum mode.

Parameters:

a (DistributedArray) – Input distributed arrays.
b (DistributedArray) – Input distributed arrays.
out (optional) – A location into which the result is stored. This option is currently not supported.

Returns:

The matrix product of the inputs.

Return type:

DistributedArray

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.mode
'sum'
>>> C.all_chunks()
{0: [array([[0, 0],
            [0, 3]]),
     array([[0, 0],
            [6, 9]])],
 1: [array([[20, 23],
            [56, 65]])],
 2: [array([[26, 29],
            [74, 83]])]}
>>> C
array([[20, 23, 26, 29],
       [56, 68, 80, 92]])