nemo_rl.environments.utils#

Module Contents#

Functions#

chunk_list_to_workers

Chunk a list into a list of lists, where each sublist is assigned to a worker. Keeps ordering of elements.

API#

nemo_rl.environments.utils.chunk_list_to_workers(
to_chunk: List[Any],
num_workers: int,
) → List[List[Any]][source]#

Chunk a list into a list of lists, where each sublist is assigned to a worker. Keeps ordering of elements.

If the list is not divisible by the number of workers, the last worker may have fewer elements. If there are more workers than elements, the first len(list) workers will have a single element each, and the remaining workers will have empty lists.

Parameters:

  • list – The list to be chunked.

  • num_workers – The number of workers to distribute the list to.

Returns:

A list of lists, where each sublist contains elements assigned to a worker.

Examples:

>>> from nemo_rl.environments.utils import chunk_list_to_workers
>>> chunk_list_to_workers([1, 2, 3, 4, 5], 3)
[[1, 2], [3, 4], [5]]