def _filter_features(
    example,
    feature_whitelist):
  """Remove features that are not whitelisted.

  Args:
    example: Input example.
    feature_whitelist: A list of feature names to whitelist.

  Returns:
    An example containing only the whitelisted features of the input example.
  """
  return {
      feature_name: example[feature_name]
      for feature_name in feature_whitelist
      if feature_name in example
  }