Factor out some Operation mutations from Function into Operation.

These mutations are used to rewrite a forward function call to output its
intermediates when creating a gradient call. This change is in preparation
for making cond_v2 and while_v2 do the same thing, instead of outputing
all intermediates from the forward call when its created.

PiperOrigin-RevId: 223250964