Class NDArrayPartition

python/dgl/partition.py:474–632

Create a new partition of an NDArray. That is, an object which assigns each row of an NDArray to a specific partition. Parameters ---------- array_size : int The first dimension of the array being partitioned. num_parts : int The number of parts to divide the arr

class NDArrayPartition(object):
    """Create a new partition of an NDArray. That is, an object which assigns
    each row of an NDArray to a specific partition.

    Parameters
    ----------
    array_size : int
        The first dimension of the array being partitioned.
    num_parts : int
        The number of parts to divide the array into.
    mode : String
        The type of partition. Currently, the only valid values are
        'remainder' and 'range'.
        'remainder' assigns rows based on remainder when dividing the row id by the
        number of parts (e.g., i % num_parts).
        'range' assigns rows based on which part of the range 'part_ranges'
        they fall into.
    part_ranges : Tensor or dgl.NDArray, Optional
        Should only be specified when the mode is 'range'. Should be of the
        length `num_parts + 1`, and be the exclusive prefix-sum of the number
        of nodes in each partition. That is, for 3 partitions, we could have
        the list [0, a, b, 'array_size'], and all rows with index less
        than 'a' are assigned to partition 0, all rows with index greater than
        or equal to 'a' and less than 'b' are in partition 1, and all rows
        with index greater or equal to 'b' are in partition 2. Should have
        the same context as the partitioned NDArray (i.e., be on the same GPU).

    Examples
    --------

    A partition of a homgeonous graph `g`, where the vertices are
    striped across processes can be generated via:

    >>> from dgl.partition import NDArrayPartition
    >>> part = NDArrayPartition(g.num_nodes(), num_parts, mode='remainder' )

    A range based partition of a homogenous graph `g`'s nodes, where
    the nodes are stored in contiguous memory. This converts an existing
    range based partitioning (e.g. from a
    dgl.distributed.graph_partition_book.RangePartitionBook)
    'max_node_map', to an NDArrayPartition 'part'.

    >>> part_range = [0]
    >>> for part in part_book.metadata():
    >>>     part_range.append(part_range[-1] + part['num_nodes'])
    >>> part = NDArrayPartition(g.num_nodes(), num_parts, mode='range',
    ...                         part_ranges=part_range)
    """

    def __init__(
        self, array_size, num_parts, mode="remainder", part_ranges=None
    ):
        assert num_parts > 0, 'Invalid "num_parts", must be > 0.'
        if mode == "remainder":
            assert part_ranges is None, (
                "When using remainder-based "
                'partitioning, "part_ranges" should not be specified.'
            )