gigl.distributed.utils.get_ids_on_rank

gigl.distributed.utils.partition_book.get_ids_on_rank(partition_book: Tensor | PartitionBook, rank: int) → Tensor

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. Args:

partition_book (Union[torch.Tensor, PartitionBook]): Tensor or range-based partition book rank (int): Rank of current machine