def extract_matching_values(
    x: Union[dict, list], predicate: Callable
) -> Tuple[Union[dict, list], Union[dict, list]]:
    """ Return matching and nonmatching values. Keeps hierarchy. """
    if isinstance(x, dict):
        matching_vals = {}
        nonmatching_vals = {}
        for k, v in x.items():
            if isinstance(v, (list, dict)):
                match, nonmatch = extract_matching_values(v, predicate)
                if match:
                    matching_vals[k] = match
                if nonmatch or not v:
                    nonmatching_vals[k] = nonmatch
            elif predicate(v):
                matching_vals[k] = v
            else:
                nonmatching_vals[k] = v
    else:
        assert isinstance(x, list)
        matching_vals = []
        nonmatching_vals = []
        for v in x:
            if isinstance(v, (list, dict)) and v:
                match, nonmatch = extract_matching_values(v, predicate)
                if match:
                    matching_vals.append(match)
                if nonmatch or not v:
                    nonmatching_vals.append(nonmatch)
            elif predicate(v):
                matching_vals.append(v)
            else:
                nonmatching_vals.append(v)
    return matching_vals, nonmatching_vals
 Return matching and nonmatching values. Keeps hierarchy. 
training/DeepSpeed-Domino/megatron/core/dist_checkpointing/dict_utils.py