core.dist_checkpointing.dict_utils#

Utilities for operating with dicts and lists.

All functions in this module handle nesting of dicts and lists. Other objects (e.g. tuples) are treated as atomic leaf types that cannot be traversed.

Module Contents#

Functions#

extract_matching_values

Return matching and nonmatching values. Keeps hierarchy.

diff

Recursive diff of dicts.

inspect_types

Helper to print types of (nested) dict values.

nested_values

Returns iterator over (nested) values of a given dict or list.

nested_items_iter

Returns iterator over (nested) tuples (container, key, value) of a given dict or list.

dict_map

map equivalent for dicts.

dict_map_with_key

map equivalent for dicts with a function that accepts tuple (key, value).

dict_list_map_inplace

Maps dicts and lists in-place with a given function.

dict_list_map_outplace

Maps dicts and lists out-of-place with a given function.

merge

Merges dicts and lists recursively.

map_reduce

Simple map-reduce implementation following more_itertools.map_reduce interface.

API#

core.dist_checkpointing.dict_utils.extract_matching_values(
x: Union[dict, list],
predicate: Callable[[Any], bool],
return_lists_as_dicts: bool = False,
) → Tuple[Union[dict, list], Union[dict, list]]#

Return matching and nonmatching values. Keeps hierarchy.

Parameters:

  • x (Union[dict, list]) – state dict to process. Top-level argument must be a dict or list

  • predicate (object -> bool) – determines matching values

  • return_lists_as_dicts (bool) – if True, matching lists will be turned into dicts, with keys indicating the indices of original elements. Useful for reconstructing the original hierarchy.

core.dist_checkpointing.dict_utils.diff(
x1: Any,
x2: Any,
prefix: Tuple = (),
) → Tuple[list, list, list]#

Recursive diff of dicts.

Parameters:

  • x1 (object) – left dict

  • x2 (object) – right dict

  • prefix (tuple) – tracks recursive calls. Used for reporting differing keys.

Returns:

tuple of: - only_left: Prefixes present only in left dict - only_right: Prefixes present only in right dict - mismatch: values present in both dicts but not equal across dicts. For tensors equality of all elems is checked. Each element is a tuple (prefix, type of left value, type of right value).

Return type:

Tuple[list, list, list]

core.dist_checkpointing.dict_utils.inspect_types(
x: Any,
prefix: Tuple = (),
indent: int = 4,
)#

Helper to print types of (nested) dict values.

core.dist_checkpointing.dict_utils.nested_values(x: Union[dict, list])#

Returns iterator over (nested) values of a given dict or list.

core.dist_checkpointing.dict_utils.nested_items_iter(x: Union[dict, list])#

Returns iterator over (nested) tuples (container, key, value) of a given dict or list.

core.dist_checkpointing.dict_utils.dict_map(f: Callable, d: dict)#

map equivalent for dicts.

core.dist_checkpointing.dict_utils.dict_map_with_key(f: Callable, d: dict)#

map equivalent for dicts with a function that accepts tuple (key, value).

core.dist_checkpointing.dict_utils.dict_list_map_inplace(
f: Callable[[core.dist_checkpointing.dict_utils.U], core.dist_checkpointing.dict_utils.V],
x: Union[Dict, List, core.dist_checkpointing.dict_utils.U],
)#

Maps dicts and lists in-place with a given function.

core.dist_checkpointing.dict_utils.dict_list_map_outplace(
f: Callable[[core.dist_checkpointing.dict_utils.U], core.dist_checkpointing.dict_utils.V],
x: Union[Dict, List, core.dist_checkpointing.dict_utils.U],
) → Union[Dict, List, core.dist_checkpointing.dict_utils.V]#

Maps dicts and lists out-of-place with a given function.

core.dist_checkpointing.dict_utils.merge(
x1: Union[dict, list],
x2: Union[dict, list],
key: Tuple[Union[str, int], ...] = (),
)#

Merges dicts and lists recursively.

core.dist_checkpointing.dict_utils.map_reduce(
xs: Iterable,
key_fn: Callable = lambda x: ...,
value_fn: Callable = lambda x: ...,
reduce_fn: Callable = lambda x: ...,
) → dict#

Simple map-reduce implementation following more_itertools.map_reduce interface.