cupyx.distributed.array.make_2d_index_map

cupyx.distributed.array.make_2d_index_map(i_partitions, j_partitions, devices)

为具有指定块的二维矩阵创建一个 index_map。

参数:

• i_partitions (list of ints) – i 轴上的块边界

• j_partitions (list of ints) – j 轴上的块边界

• devices (2D list of sets of ints) – 拥有每个块的设备

返回:

index_map 设备指定ID将拥有的块的索引。

返回类型:

dict from int to array indices

示例

>>> index_map = make_2d_index_map(
...     [0, 2, 4], [0, 3, 5],
...     [[{0}, {1}],
...      [{2}, {0, 1}]])
>>> pprint(index_map)
{0: [(slice(0, 2, None), slice(0, 3, None)),
     (slice(2, 4, None), slice(3, 5, None))],
 1: [(slice(0, 2, None), slice(3, 5, None)),
     (slice(2, 4, None), slice(3, 5, None))],
 2: [(slice(2, 4, None), slice(0, 3, None))]}