def GetFromNestedMapOrDie(nested_map, key):
  """Returns the value of `key` if it exists, otherwise dies with a nice error.

  Args:
    nested_map: The input `NestedMap`.
    key: The key to retrieve. May use dots to refer to a `nested.field_name`.
  """
  value = nested_map.Get(key)
  if value is None:
    all_keys = [k for k, _ in nested_map.FlattenItems()]
    raise ValueError(
        'No feature {} in input batch. Available features are: {}'.format(
            key, all_keys))
  return value