code: def map_reduce(
      xs: Iterable,
      key_fn: Callable = lambda x: x,
      value_fn: Callable = lambda x: x,
      reduce_fn: Callable = lambda x: x,
  ) -> dict:
      res = defaultdict(list)
      for x in xs:
          res[key_fn(x)].append(value_fn(x))
      for k in res:
          res[k] = reduce_fn(res[k])
      return dict(res)
  
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 12 for calling a global function with example usage: map_reduce(xs, key_fn, value_fn, reduce_fn) and returns: dict