from typing import List
import torch


def get_gpus_for_rank(world_size: int) -> List[List[int]]:
    """This will return a list, each element of which contains a list of GPUs
    to be used by the respective process.

    Examples (results are shown for a machine with 2 GPUs):

        >>> get_gpus_for_rank(2)  # [[0], [1]]
        >>> get_gpus_for_rank(4)  # [[0], [0], [1], [1]]
        >>> get_gpus_for_rank(1)  # [[0, 1]]

    Args:
        world_size (int): denotes number of subsets to split the available GPUs into
    """

    visible_devices = list(range(torch.cuda.device_count()))
    num_visible_devices = torch.cuda.device_count()

    if num_visible_devices >= world_size:
        gpus_for_rank = [[i] for i in range(world_size)]
    else:
        visible_devices_repeated = [
            [device]
            for device in visible_devices
            for _ in range((world_size + num_visible_devices - 1) // num_visible_devices)
        ]
        gpus_for_rank = visible_devices_repeated[:world_size]

    return gpus_for_rank