Add capability to forward some features to predictions dictionary in Estimator.

From @rhaertel80:
There are cases where it is useful to have externally defined keys and have these keys passed through from the input to the output. As an example, consider a batch prediction service: The service simply runs inference on the users graph and returns the results. Keys are essential because there is no order guarantee on the outputs so they need  to be rejoined to the inputs via keys or transclusion of the inputs in the outputs.

PiperOrigin-RevId: 170521852