code: def _flatten_dict(original_dict):
    
    flat_dict = {}
    for key, value in original_dict.items():
      if isinstance(value, dict):
        for name, tensor in value.items():
          if isinstance(tensor, dict):
            raise ValueError("flatten_dict only handles 2 levels of nesting.")
          flat_key = "__" + key + "_" + name
          flat_dict[flat_key] = tensor
      else:
        flat_dict[key] = value
  
    return flat_dict
docs: Flatten dict of dicts into a single dict with appropriate prefixes.
  
    Handles only 2 levels of nesting in the original dict.
  
    Args:
      original_dict: Dict which may contain one or more dicts.
    Returns:
      flat_dict: Dict without any nesting. Any dicts in the original dict have
        their keys as prefixes in the new dict.
    Raises:
      ValueError if the original dict has more than two levels of nesting.