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

  Args:
    batch: A dict containing the input batch of examples.
    feature_whitelist: A list of feature names to whitelist.

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