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.

def _flatten_dict(original_dict):
  """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.
  """
  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