def collect_required_arguments(req_args, all_kwargs):
  """Extract a dictionary with the required keys from a larger dictionary.

  Args:
    req_args: List of keys to extract
    all_kwargs: Dictionary with a superset of the required keys.
  Returns:
    req_args_dict: Dictionary with all the required arguments.
  """
  return {key: all_kwargs[key] for key in req_args}