composer.utils.iter_helpers#

Utilities for iterating over collections.

Functions

ensure_tuple

Converts x into a tuple.

iterate_with_pbar

Iterate over a batch iterator and update a tqdm.tqdm progress bar by the batch size on each step.

map_collection

Apply map_fn on each element in collection.
composer.utils.iter_helpers.ensure_tuple(x)[source]#

Converts x into a tuple.

  • If x is None, then tuple() is returned.

  • If x is a tuple, then x is returned as-is.

  • If x is a list, then tuple(x) is returned.

  • If x is a dict, then tuple(v for v in x.values()) is returned.

Otherwise, a single element tuple of (x,) is returned.

Parameters

x (Any) โ€“ The input to convert into a tuple.

Returns

tuple โ€“ A tuple of x.

composer.utils.iter_helpers.iterate_with_pbar(iterator, progress_bar=None)[source]#

Iterate over a batch iterator and update a tqdm.tqdm progress bar by the batch size on each step.

This function iterates over iterator, which is expected to yield batches of elements. On each step, the batch is yielded back to the caller, and the progress_bar is updated by the length of each batch.

Note

It is expected that the progress_bar = tqdm.tqdm(total=sum(len(x) for x in iterator)).

Parameters

  • iterator (Iterator[TSized]) โ€“ An iterator that yields batches of elements.

  • progress_bar (Optional[tqdm.tqdm], optional) โ€“ A tqdm.tqdm progress bar. If None (the default), then this function simply yields from iterator.

Yields

Iterator[TSized] โ€“ The elements of iterator.

composer.utils.iter_helpers.map_collection(collection, map_fn)[source]#

Apply map_fn on each element in collection.

  • If collection is a tuple or list of elements, map_fn is applied on each element, and a tuple or list, respectively, containing mapped values is returned.

  • If collection is a dictionary, map_fn is applied on each value, and a dictionary containing the mapped values is returned.

  • If collection is None, None is returned.

  • If collection is a single element, the result of applying map_fn on it is returned.

Parameters

  • collection โ€“ The element, or a tuple of elements.

  • map_fn โ€“ A function to invoke on each element.

Returns

  • Collection โ€“ The result of applying map_fn on each element of collection.

  • The type of ``collection`` is preserved.