[tf.data] Fix Python shape inference for `tf.contrib.data.map_and_batch()`.

Previously, it would incorrectly report that all batches have the same size, not accounting for the possibility of the last batch being partial.

Fixes #17720.

PiperOrigin-RevId: 189121488