cupyx.distributed.array.make_2d_index_map#

cupyx.distributed.array.make_2d_index_map(i_partitions, j_partitions, devices)[source]#

Create an index_map for a 2D matrix with a specified blocking.

Parameters:

  • i_partitions (list of ints) – boundaries of blocks on the i axis

  • j_partitions (list of ints) – boundaries of blocks on the j axis

  • devices (2D list of sets of ints) – devices owning each block

Returns:

index_map

Indices for the chunks that devices with designated IDs are going to own.

Return type:

dict from int to array indices

Example

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