cupyx.distributed.array.make_2d_index_map

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

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

参数:

• i_partitions (列表 of 整数) – i 轴上块的边界

• j_partitions (列表 of 整数) – j 轴上块的边界

• devices (由整数集合组成的二维列表) – 拥有每个块的设备

返回:

index_map

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

返回类型:

从整数到数组索引的字典

示例

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