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.

Module Contents#

gigl.distributed.utils.partition_book.get_ids_on_rank(partition_book, rank)[source]#

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