def flatten_nested_dict(x: Dict[str, Any],
                        join_str: str = '/',
                        prefix: str = '') -> Dict[str, Any]:
  """Transforms nested dictionary into a flat dictionary."""
  assert isinstance(x, dict)
  result = {}
  for k, v in x.items():
    key = prefix + join_str + k
    if isinstance(v, dict):
      result.update(flatten_nested_dict(v, join_str, key))
    else:
      result[key] = v
  return result
Transforms nested dictionary into a flat dictionary.
language/mentionmemory/utils/checkpoint_utils.py