Rewrite the postprocess logic to support dynamic batch size better.
Rewrite the postprocess logic to support dynamic batch size better.

PiperOrigin-RevId: 338732758

tensorflow/tpu,tensorflow/tpu,tensorflow/tpu,tensorflow/tpu