Small cleanup for partial_eval (#3210)
`partial_eval` uses some pretty tricky conventions for return values
(see `partial_eval_wrapper`), but it forces all call sites to deal with
untangling them. This commit inlines the postprocessing into
`partial_eval`, greatly simplifying its usage.
