code: def 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.
  
      Arguments:
          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.
      """
  
      def _set_elem(target, k, v):
          if return_lists_as_dicts:
              target[k] = v
          else:
              target.append(v)
  
      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, return_lists_as_dicts
                  )
                  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
      elif isinstance(x, list):
          matching_vals = {} if return_lists_as_dicts else []
          nonmatching_vals = {} if return_lists_as_dicts else []
          for ind, v in enumerate(x):
              if isinstance(v, (list, dict)) and v:
                  match, nonmatch = extract_matching_values(
                      v, predicate, return_lists_as_dicts
                  )
                  if match:
                      _set_elem(matching_vals, ind, match)
                  if nonmatch or not v:
                      _set_elem(nonmatching_vals, ind, nonmatch)
              else:
                  target = (
                      matching_vals
                      if predicate(v)
                      else nonmatching_vals
                  )
                  _set_elem(target, ind, v)
      else:
          raise ValueError(
              f"Unexpected top-level object type: {type(x)}"
          )
      return matching_vals, nonmatching_vals
  
file_path: megatron-lm/megatron/core/dist_checkpointing/dict_utils.py
imports: from collections import defaultdict
  from typing import Any, Callable, Iterable, Optional, Tuple, Union
  import torch
label_desc: Use image node_id 1 for calling a global function with example usage: extract_matching_values(x, predicate, return_lists_as_dicts) and returns: matching_vals,  nonmatching_vals