gigl.distributed.utils.partition_book

Functions

get_ids_on_rank(partition_book, rank)	Provided a tensor-based partition book or a range-based bartition book and a rank, returns all the ids that are stored on that rank.
get_total_ids(partition_book)	Returns the total number of ids (e.g. the total number of nodes) from a partition book.

Module Contents

gigl.distributed.utils.partition_book.get_ids_on_rank(partition_book, rank)

Provided a tensor-based partition book or a range-based bartition book and a rank, returns all the ids that are stored on that rank. :param partition_book: Tensor or range-based partition book :type partition_book: Union[torch.Tensor, PartitionBook] :param rank: Rank of current machine :type rank: int

Parameters:

• partition_book (Union[torch.Tensor, graphlearn_torch.partition.PartitionBook])

• rank (int)

Return type:

torch.Tensor

gigl.distributed.utils.partition_book.get_total_ids(partition_book)

Returns the total number of ids (e.g. the total number of nodes) from a partition book. :param partition_book: Tensor or range-based partition book :type partition_book: Union[torch.Tensor, PartitionBook]

Returns:

Total number of ids in the partition book

Return type:

int

Parameters:

partition_book (Union[torch.Tensor, graphlearn_torch.partition.PartitionBook])