def nest_to_flat_dict(nest):
  """Convert a nested structure into a flat dictionary.

  Args:
    nest: A nested structure.

  Returns:
    flat_dict: A dictionary with strings keys that can be converted back into
               the original structure via `flat_dict_to_nest`.
  """
  flat_sequence = tf.nest.flatten(nest)
  return {str(k): v for k, v in enumerate(flat_sequence)}
Convert a nested structure into a flat dictionary.

  Args:
    nest: A nested structure.

  Returns:
    flat_dict: A dictionary with strings keys that can be converted back into
               the original structure via `flat_dict_to_nest`.
  
language/common/utils/nest_utils.py