dgl.metis_partition_assignment

dgl.metis_partition_assignment(g, k, balance_ntypes=None, balance_edges=False, mode='k-way')

This assigns nodes to different partitions with Metis partitioning algorithm.

When performing Metis partitioning, we can put some constraint on the partitioning. Current, it supports two constrants to balance the partitioning. By default, Metis always tries to balance the number of nodes in each partition.

• balance_ntypes balances the number of nodes of different types in each partition.

• balance_edges balances the number of edges in each partition.

To balance the node types, a user needs to pass a vector of N elements to indicate the type of each node. N is the number of nodes in the input graph.

After the partition assignment, we construct partitions.

Parameters

• g (DGLGraph) – The graph to be partitioned

• k (int) – The number of partitions.

• balance_ntypes (tensor) – Node type of each node

• balance_edges (bool) – Indicate whether to balance the edges.

• mode (str, "k-way" or "recursive") – Whether use multilevel recursive bisection or multilevel k-way paritioning.

Returns

A vector with each element that indicates the partition ID of a vertex.

Return type

a 1-D tensor