def rekey(batch, key_map):
  """Rekeys a batch using key_map.

  key_map specifies new_key: old_key pairs.

  Args:
    batch: a dictionary to modify.
    key_map: a dictionary that new keys to old keys. So if you want your new
      dataset to have keys 'inputs' and 'targets', where the old task just had
      'targets, pass {'inputs': 'targets', 'targets': 'targets'}.

  Returns:
    a new batch dict.
  """
  return {key: batch[value] for key, value in key_map.items()}