import functools


def with_post_processing(fn, post_process_fn):
  """Generates a function that applies `post_process_fn` to outputs of `fn`."""
  @functools.wraps(fn)
  def apply_fn(*args, **kwargs):
    return post_process_fn(*fn(*args, **kwargs))

  return apply_fn