Add an output context that can be used to specify outputs to capture when...

Add an output context that can be used to specify outputs to capture when running multiple steps at a time using the `run_steps_on_dataset` API. It allows the user's step function to specify which outputs to emit at what frequency. Currently it only supports capturing output from the last step, but will soon be augmented to support other use cases such as output each N steps.

PiperOrigin-RevId: 202520245