cupyx.distributed.array.matmul#

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

分布式数组之间的矩阵乘法。

参数必须具有兼容的 shape 和 index_map。

此操作将其操作数转换为副本模式 (replica mode)，并在求和模式 (sum mode) 下计算它们的乘积。

参数:

a (DistributedArray) – 输入分布式数组。
b (DistributedArray) – 输入分布式数组。
out (可选) – 存储结果的位置。此选项目前不支持。

返回:

输入的矩阵乘积。

返回类型:

DistributedArray

示例

>>> 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]])

另请参阅

numpy.matmul