code: def dict_map_with_key(f: Callable, d: dict):
      for sub_d, k, v in nested_items_iter(d):
          sub_d[k] = f(k, v)
  
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 8 for calling a global function with example usage: dict_map_with_key(f, d) without return types